TL;DR. A World Action Model (WAM) is a robot policy that learns to predict its next actions and what that action will do. I built a small one, pico-WAM (1.3M parameters, about 1k lines of PyTorch), trained it on the push-T dataset, and checked every prediction against the simulator. This blog gives a thorough background about how WAMs are built, across models and how pico-WAM is structured.

Terminologies

Here are some initial ideas that will help you throughout this tutorial.

1. A policy is a function that generates next-action. You give it what the robot sees (a camera image), what it knows about itself and the scene (positions and velocities, the state), and what you want (an instruction like “push the T onto the target”). It gives back an action.
2. Imitation learning (behavioral cloning, BC) is the simplest way to train a policy: record an expert doing the task, then train the network to copy the expert’s action in every situation. It is plain supervised learning, inputs to labels. A VLA (vision-language-action model) is a policy trained this way that takes images and language as input.
3. A world model (or forward model) answers a different question: if I am here and I do this, where will I be next? Engineers usually write this by hand, that’s why we have physics simulators and model-predictive control (MPC). A learned world model learns this from data instead.
4. A latent is a compact vector the network uses internally to summarize what it sees. You can’t read it directly, but you can train small “readout” networks to turn it back into things you understand: positions, velocities, pictures.
5. Imagination, or a rollout, means applying the world model again and again without looking at the real world. If the predictions stay good for many steps, the model has learned something real about the world.

That’s it! A WAM combines these ideas in one network.

VLA vs World-Model vs WAM

Terminologies around WAMs are quite confusing, so we will logically define our own scope. A recent survey defines WAMs as models that “connect future-state modeling with executable action generation” (Lu et al.). Another sorts them by the future they produce: pixels, a latent, or nothing at all (Shen et al.). When I was getting started with WAMs, I was constantly juggling around what these terms mean in this context, and turns out, most of it is something I have learned before with a different name.

Designing a WAM

What is predicted? You can predict raw future frames (expensive, but you can look at them), a compressed learned representation (cheap, but harder to inspect), the physical state (only possible when you know it, as in simulation), or nothing at all. pico-WAM predicts the physical state, and optionally pixels.

How do actions meet the future? You can generate the action and the future together (“joint prediction”). You can generate the future first, then ask an inverse dynamics model which action gets you there (“plan, then act”). Or you can build an explicit forward model that takes any action as input. pico-WAM trains jointly and has an explicit action-conditioned forward model, which is what makes “what if” questions and planning possible.

Which backbone? Big systems reuse pretrained video-generation models. I use small MLPs so everything trains on a laptop.

What is the prediction for? It can shape the representation during training (the cheapest option and the most common), rank candidate actions at test time (planning), produce the action directly, or generate fake training data for another policy.

Open-source WAMs

(full disclosure: LLM helped me write this section because I constantly kept missing some important takeaways that distinguish these models from each other… so, live long ai!)

• GR-1 (Wu et al.) is a GPT-style transformer, pretrained to predict future frames from Ego4D human videos, then fine-tuned to output actions and future images. The future prediction is an auxiliary training signal.
• UWM (Zhu et al.) puts video diffusion and action diffusion into one transformer, with a separate noise level for each. By choosing which one to denoise, the same network acts as a policy, a forward model, an inverse model or a video generator. That trick also lets it learn from videos that have no action labels.
• WorldVLA (Cen et al.) turns images, text and actions into discrete tokens in one autoregressive multimodal LLM, which learns both to output action tokens and to predict the next frame.
• Genie Envisioner (Liao et al.) is a 2B-parameter video diffusion transformer pretrained on about a million real robot episodes, with a lightweight flow-matching action module that reads its video latents. The video model doubles as a neural simulator.
• Cosmos Policy (Kim et al.) fine-tunes NVIDIA’s 2B Cosmos-Predict2 video model without changing its architecture: actions, robot state and value estimates are written into the video latent as extra “frames”. At test time it plans by sampling several action proposals and picking the one whose predicted future looks most valuable.
• V-JEPA 2-AC (Assran et al.) never predicts pixels. It trains a 300M-parameter action-conditioned predictor on top of a frozen self-supervised video encoder, and plans with model-predictive control in that representation space. It has no action head at all; the planner picks the actions.
• DreamGen (Jang et al.) uses a video world model as a data generator. It fine-tunes a video model on a few robot demos, generates many new videos, labels them with an inverse-dynamics or latent-action model, and trains an ordinary policy on the result.
• LingBot-VA (Li et al.) is a 5.3B-parameter autoregressive video-action model that predicts future video latents and reads actions off them with inverse dynamics. Its KV cache keeps a long memory of the past.

Two systems that plan on Push-T itself are worth knowing too. DINO-WM (Zhou et al.) plans with the cross-entropy method in the features of a frozen DINOv2 encoder, with a predictor of about 19M parameters trained on 18,500 Push-T trajectories. PILOT (Liu et al.) starts that kind of search from a learned policy’s actions and beats executing the same policy on long Push-T tasks.

Outro

All this said, I hope now you have a better understanding of WAMs. In the next blog, we will go deeper into each layer of pico-WAM.

The Dream of a Generalist Robot

The goal of a generalist robot is to move beyond the factory floor, where machines repeat a single task millions of times. We want robots that can navigate messy, unpredictable human environments, such as cleaning up a kitchen or folding laundry.

The Generalization Problem (again)

Why has this been so hard? The core obstacle is generalization. Traditional robots are “brittle”; if you train them to pick up a red block in one spot, they fail if you swap it for a blue ball or move it two inches to the left. They learn specific choreography rather than the underlying concept of an action.

Predicting the World, Instead of Just Actions

DreamZero introduces the World Action Model (WAM). Instead of simply mapping “If I see this, do that,” the robot learns to predict the future. By “dreaming” or generating video clips of how the world evolves based on its actions, the robot builds a deep intuition for physics and cause-and-effect.

Physics in WAMs

WAM treats the physical world as something to be predicted. Let me give you an intuition.

Predictive Modeling

Unlike traditional robots that simply react to inputs, DreamZero asks, given the current state and my intended actions, what will the world look like in a few moments? It learns to dream (aka simulate future video frames) as a representation of cause and effect.

Autoregressive Diffusion

It uses a massive 14-billion parameter autoregressive diffusion transformer. The diffusion process allows it to denoise potential futures, sculpting coherent physical sequences out of static, while the autoregressive nature ensures that each movement logically flows from the last.

Real-World Grounding

To ensure its “dreamed” physics match reality, the robot follows a continuous loop: Observe, Predict, Act, and Update. By constantly comparing its predictions with actual camera observations, it corrects its internal model, grounding its imagination in the actual laws of the environment.

How DreamZero Works

At its core, DreamZero is a 14-billion parameter autoregressive diffusion transformer.

• Diffusion. It denoises static, “sculpting” future video frames and actions out of a noisy, predicted state.
• Autoregressive. Like a large language model, it generates sequences where each movement logically flows from the last, ensuring fluid, lifelike motion rather than jerky, disconnected steps

Dream to Reality Transfer

How does it stay accurate? The robot operates in a four-step loop: Observe, Predict, Act, and Update. The secret-sauce is the update step. By constantly comparing its dreamed predictions against the actual real-world frames from its cameras, the robot corrects its course, ensuring its imagination remains grounded in reality.

What’s Next?

DreamZero is currently a “System 1” thinker — excellent at fast, reactive tasks but limited by a 6-second visual memory. The next major frontier is “System 2” deliberate, long-term planning. As researchers scale these video world models, we are likely to see an explosion in robotic capabilities, changing the way machines interact with our world forever.

Epilogue

I hope you have a better idea of how WAMs solve some physics-related concepts better than VLAs. I will be posting follow-up blogs as I learn about WAMs in more details (super excited to share more insights that are in progress right now) which will dive into technical depths of these concepts alongside more references.

I’ve been thinking about this a lot lately, partly because of an article I came across — Business Insider’s piece on robots coming for the jobs no one wants. The framing is familiar: aging workforce, labor shortages in repetitive tasks, companies turning to humanoid automation to fill the gap. And it’s not wrong. Firms like Agility Robotics are deploying bipedal robots into real manufacturing and logistics facilities. The pitch is compelling — these humanoids slot into existing factory layouts without ripping up the floor plan, and they’re built for the kind of mundane physical labor that fewer humans want to do.

What’s accelerating this is the infrastructure layer. New physical AI models from NVIDIA are giving these machines something closer to a “digital brain” — the ability to actually perceive and navigate a messy, unstructured world. As someone who’s spent time training VLA models and watching them struggle with even basic spatial reasoning, I can tell you: the jump from “works in simulation” to “works on a real table with real lighting and real object clutter” is enormous. So when I see companies shipping robots that can handle even a narrow slice of real-world tasks reliably, I pay attention.

And the impact is going global faster than I expected. In Europe, in China, humanoid robots are beginning to assemble cars alongside human workers. These aren’t lab prototypes anymore — they’re on the production line. But here’s where my own work makes me pump the brakes a little.

Rodney Brooks made a point in Why Today’s Humanoids Won’t Learn Dexterity that I think deserves more airtime than it gets: despite these robots’ impressive forms, they still lack the fine-motor adaptability that makes human manipulation so fluid. I feel this tension every day.

So there’s a genuine tension here: we’re deploying robots to save supply chains and fill labor gaps, and that’s working for structured, repetitive tasks. But the dexterous, adaptive manipulation that makes human labor truly versatile? We’re years away from cracking that — and honestly, anyone who tells you otherwise is probably fundraising.

Which brings me to what actually bothers me about the current hype cycle. Open your social feed and the revolution looks like it already happened. A bipedal robot loads car parts at a BMW plant. A drone drops a medical package into a remote village. A warehouse arm sorts a thousand parcels without blinking. The clips go viral, the comments erupt, another round of funding gets announced at a valuation that makes your head spin — Figure AI closed a Series C north of $1 billion at a reported $39 billion valuation.

But strip away the cinematography, and the question that actually matters is deeply unglamorous: can it do that again tomorrow? And the day after? And the day after that — without someone flying in to fix it?

I’ve watched enough “it works!” moments in the lab dissolve into “wait, why did it stop working?” to know that repeatability is the real product.

That’s the unsexy frontier. And it’s the one I think we should be talking about more.

Numbers Behind the Noise

Before we talk about where the exciting stuff is happening, it helps to ground ourselves in what “deployment” actually looks like at scale — because the bar is higher than most demo reels suggest.

The IFR reported 542,000 industrial robots installed in 2024. That’s more than double what it was a decade ago. Asia accounts for 74% of new deployments, with Europe at 16% and the Americas trailing at 9%. On the service robotics side, professional units hit around 199,000 in 2024 — up 9% — with transportation and logistics leading the pack at nearly 103,000 units. Hospitality came in around 42,000, and cleaning robots saw the sharpest growth at over 25,000 units, climbing 34%.

Here’s the thing worth noting: almost none of those numbers are about humanoids. This is the baseline — the boring, proven, already-profitable version of “robots in the real world.” When a warehouse operator evaluates a new system, they’re comparing it against this track record of reliability and ROI, not against a viral demo. That’s the buyer expectation any new form factor has to meet.

Where Robots Are Actually Earning Their Keep

The “demos → deployments” transition is clearest in places where the environment is structured, the task is repetitive, and uptime directly translates to money — warehouses, factories, hospitals, logistics networks. The IFR data backs this up: the biggest professional service robot category by units shipped is indoor transport and warehouse logistics. Industrial installations stay concentrated in Asia, where manufacturing scale and automation intensity create the economics that justify the investment.

The clearest example of deployment-as-product might be Amazon. They’ve deployed over a million robots across their operations network since 2012 and are layering AI systems on top to optimize fleet movement. What makes this interesting isn’t the robot count — it’s that each robot is only valuable as part of a continuously tuned system. The facilities, the processes, the safety constraints, the software updates, the maintenance schedules, the human operator training — take any one of those away and the robot becomes expensive furniture. That’s the part most people skip when they share the headline number.

The Tech Stack That’s Actually Enabling This

So what’s making real deployment possible right now? A few threads are converging, and they’re worth unpacking because they map directly to where I think the field is headed.

Data and generalist policy training. The Open X-Embodiment effort was a big deal — a standardized multi-robot dataset with 1M+ real trajectories spanning 22 different robot embodiments, plus models that demonstrate genuine cross-embodiment transfer. It reduces the cost of adding new skills across fleets and platforms. Why does that matter practically? Because it chips away at one of the most painful bottlenecks in robotics: every time you want a new skill on a new platform, you used to be starting almost from scratch. If you can amortize training data across fleets and form factors, the economics of adding capabilities shift dramatically. I’ve been watching this closely since building Nemotron-VLA — the promise of generalist models is real, but so is the gap between “transfers in evaluation” and “transfers reliably in production.”

VLAs and embodied reasoning as integration layer. Google DeepMind’s Gemini Robotics 1.5 technical report frames something I think is genuinely important: multi-embodiment VLA models that fuse perception + language + action with physical reasoning. The idea of a model that can see a scene, understand a natural language instruction, and plan a physically grounded action sequence. I’ve spent enough time training and debugging VLA architectures to know we’re not there yet, but the trajectory is clear. The question is less “will this work?” and more “how narrow does the task domain need to be before it works reliably?”

Touch and force sensing as a practical differentiator. This one is close to my heart. Brooks has been insisting for years that real dexterity depends on touch, and he’s being proven right — not by humanoids suddenly mastering human-like hands, but by targeted industrial systems that actually instrument tactile feedback. Amazon’s Vulcan is a great example: it’s their first robot with what they call a “sense of touch,” using force/torque sensing to manage contact-rich manipulation. It reportedly handles about 75% of items in a typical fulfillment inventory and has already processed serious order volumes in real deployments. That 75% number is telling — it’s impressive for a deployed system, but it also means one in four items still needs a human. That’s the kind of honest metric I wish more companies would publish.

Hybrid control pipelines: classical + learning + simulation. This is the one I feel strongest about, because I’ve lived it. The deployment reality for high-value robots is almost always hybrid — you blend learning-from-demonstration, RL fine-tuned in simulation, and good old-fashioned classical control to meet the safety and reliability bar that real facilities demand. No single paradigm gets you there alone.

Pure RL is too brittle for safety-critical contact. Pure classical control can’t handle the variability of real environments.

But layer them together thoughtfully — use the learned policy for the adaptive parts, the classical controller for the safety-critical parts, simulation for the edge cases you can’t afford to hit in the real world — and you get something that actually works. I’ve seen teams adopt this approach, and I’ve worked alongside some of them. It works like a charm most of the time, and when it doesn’t, you at least know which layer broke, which is more than you can say for an end-to-end black box.

So Where Does That Leave Us?

I spend my days in a lab at Northeastern working on dexterous manipulation and reinforcement learning. I’ve built VLA models from scratch, trained diffusion policies on real robot data, and watched a gripper I designed in MuJoCo struggle with tasks a toddler does without thinking. I’ve also spent time at Familiar Machines working across the full stack — AI, hardware, software — and taught a course on RL where I had to explain to students why the gap between a beautiful Bellman equation on a whiteboard and a robot that doesn’t drop things is so maddeningly wide.

All of which is to say: I’m not a spectator in this conversation. And from where I’m standing, here’s what I actually believe.

The revolution is real, but it’s slower and messier than the timeline X/LinkedIn wants to sell you. The robots that are changing the world right now aren’t the ones going viral — they’re the unglamorous fleet vehicles in Amazon warehouses, the AGVs quietly moving pallets in Shenzhen, the surgical systems doing the same procedure for the thousandth time with sub-millimeter consistency. They win not because they’re impressive, but because they’re boring. Boring means repeatable. Repeatable means deployable. Deployable means someone signs a purchase order.

The humanoids will get there. The VLAs will get better. Touch sensing will mature. But the companies that win won’t be the ones with the best demo reel — they’ll be the ones that figured out maintenance schedules, operator training, graceful failure modes, and the thousand other mundane things that turn a prototype into a product.

If you’re building in this space, my unsolicited advice is simple: fall in love with the boring parts. The deployment pipeline, the error logging, the 2am edge case. That’s where the real engineering happens, and it’s where the real value gets created.

Keivalya is an MS Robotics student at Northeastern University and he builds VLA models, teaches reinforcement learning, and occasionally convinces robots to not drop things. You can find more of his work at keivalya.com and on RoboSpace.app.

1. Importance of Open-source AI
2. Techniques from DeepSeek
   a. Group Relative Policy Optimization
   b. <think> the “Aha!” moment <think>
   c. Pure RL!
   d. Flashlight Effect in Math vs Language
   e. Distillation
3. “with great powers…”

Let’s be honest… We live in an era where companies guard AI secrets like they are some kind of neuclear weapons. Not that I don’t realize the serious implications of AI, but sometimes it’s just too much (and it seems a little unnecessary).

These papers, articles, reports — feel more like marketing brouchers rather than scientific contributions. I mean, I see you are ‘advertizing’ your model, fair enough… What are you contributing to science? It’s hard to say “AI is the new electricity” when you know physica behind electricity is free for all. I understand everyone would have a different opinion on it, and maybe I’d change mine if I become wiser/dumber in future.

On the other hand,

… and sure, there are some more efforts in open-sourcing ‘science’, and we all LOVE it!

In this Medium I’ll focus on — how do we teach LLMs to be so *awesome*, and what that teaches us in return. My goal here is to explain some really cool training/fine-tuning methods that DeepSeek used to build the foundation of world’s #1 model (at some point).

Group Relative Policy Optimization

Traditionally, AI training relied on PPO (think: an expensive private tutor) who critiques every single sentence a student writes. This is slow, costly, and limited. DeepSeek replaced this “teacher” with a “judge” through GRPO!

In this system, the student is given a single problem and generates 16 different versions of the answer. We don’t micromanage the process; we simply judge the results. Did the code run? Was the math correct? We let the versions compete.

The best ones get a medal, the crappy ones get discarded.

Just like in our life — our inner critic is a bottleneck, and inner judge is a scaling solution. Stop agonizing over your first draft or your first thought! When faced with a high-stakes decision, generate five distinct paths. Don’t critique them as you go. Let them exist, compare their projected outcomes, and commit only to the one that survives the arena of logic.

<think> the “Aha!” Moment <think/>

DeepSeek’s researchers witnessed a startling phenomenon: the AI naturally learned that “thinking” before “speaking” led to better results. Without being explicitly programmed to pause, the model began generating internal dialogue like “Wait…” or “Let me re-calculate.”

It discovered an undeniable law: thinking time directly correlates to performance.

It taught itself to have an “aha!” moment. (let that sink in for a moment)

THINK about it! The AI naturally learnt that “double checking” is does NOT mean hesitation, but a step towards brilliance. Take a breath! Deliberately pause, because, this extra time is where the intelligence actually lives!

Pure Reinforcement Learning

If you want to get better at chess, do you constantly just read the rules and learn from past matches? How many books do you read to learn to play decent guitar?

DeepSeek-R1-Zero shows that you don’t need human examples to achieve mastery; you need Pure Reinforcement Learning. By playing against itself and following the rules, the model’s success rate in competition math skyrocketed from 15% to nearly 80%.

The idea is, there is only limited amount of knowledge generated by humans so far — it becomes a ceiling, but playing against yourself has (almost) no limit. Theory is capped by the author’s knowledge, whereas practice is limited only by your persistence. Embrace the mess, and learn the game by playing it.

Flashlight Effect

While pure logic is powerful, it can lead to “wandering randomly in a dark forest.” Without a guide, the AI occasionally spiraled into gibberish or switched languages mid-sentence. Here, the idea is to give the model a few high-quality examples as “flashlight” so that it could steer the model towards human-readable logic.

Distillation

Researchers took the massive R1 model and had it write a 800,000-example “textbook” of its own thinking process. They used this to teach a tiny, 7B-parameter model.

The result was shocking: this tiny model is nearly 6 times better on competition-level math than giants like GPT-4o. It is small enough to run on a laptop, yet it carries the wisdom of a billion-dollar training run.

You need the genius to write it, but not to read it. That is the key!

If your wisdom cannot be distilled into a simple routine, it isn’t mastery; it’s just clutter.

with great power comes great responsibility

The shift from “closed” to “open” intelligence marks a move toward the democratization of wisdom. I love how DeepSeek has shown that the most powerful tools in existence are moving from the hands of the few to the hands of the many.

If a machine can learn to think, wait, and self-correct to achieve the impossible, what is stopping you from running your own internal “R1” upgrade today?

Similar to previous blogs, this is going to have technical depth alongside intuition and code interpretation. If you are new to these blogs, I am Keivalya Pandya (you can call me Kv), and I write about robotics, ai, and sometimes life. I highly recommend reading previous blogs before getting started with this one — each of them are like 8-min read, so should be super quick!

• Building Vision-Language-Action Model from scratch
• Building VLA models from scratch — II

These fore-mentioned blogs will give you a deeper understanding of how mini-VLA was built from first-principles. I carry the same idea in this blog, hence we will look into CLIP and SigLIP in fairly shallow details and try to understand how it was integrated efficiently with the existing mini-VLA architecture. The main challenges were (1) keeping the codebase minimal without turning it into a “framework” (2) making CNN ↔ ViT swappable without touching the rest of the model, and (3) getting the preprocessing right — resize + normalize so the new encoders behave as expected.

mini-VLA is simple… each one of the text/vision/state token is a single vector (B, d_model) and are fused into a diffusion action head. That’s a beautiful design constraint because it means,

As long as the vision encoder outputs a single (B, d_model) vector, everything downstream stays unchanged.

Isn’t that awesome?! So instead of modifying fusion, diffusion, training loops, etc… we only upgrade the vision token generator. That’s it! And that’s exactly what we will cover in this blog.

GitHub - keivalya/mini-vla at vision

CLIP = Contrastive Language-Image Pretraining

Now, CLIP is a popular model developed and open-source by OpenAI, and has been beautifully explained in the original paper, code, and blog

However, I want to focus on giving you a high-level overview of what am I really doing with CLIP, what I actually use it for (and what I tend to ignore from this model).

The context is, when people hear CLIP, they usually think about image–text similarity, zero-shot classification, or prompt engineering tricks. That’s not what I’m using it for here.

In mini-VLA, CLIP is simply a strong image → vector function. That’s it.

I am NOT training CLIP, nor using the text encoder, nor am I computing contrastive losses, neither am I doing some sort of prompt matching, and I am not even changing anything downstream. I am only replacing my old CNN vision stack with a pretrained CLIP vision encoder and asking a single question:

Can CLIP give me a single, semantically rich vector per image that plugs into my existing VLA pipeline?

Spoiler alert! The answer is, yes — perfectly.

But, why CLIP fits mini-VLA so well? Recall our design constraint,

Each modality (vision / language / state) is represented as a single vector of shape (B, d_model).

CLIP’s vision encoder (which is ViT-based) already outputs a sequence of patch tokens and a special pooled representation (the CLS / pooled output). This pooled output is exactly what we want here!

image → CLIP ViT → pooled embed → projection → vision token (B, d_model)

No attention maps. No spatial grids. No feature pyramid. Just one vector in, one vector out.

The hypothesis is, CLIP must work at least better than CNN, right? With CNNs we usually think in terms of receptive fields, hierarchical features, and spatial inductive biases. CLIP is different. Here we think in terms of global context (like, what is present in this image?), patch-to-patch relationships (sometimes inferring object and relative positioning in the image), and mainly semantic alignment learned from internet-scale data!

This matters a lot for robotics. Instead of encoding “edges → textures → objects”, CLIP tends to encode object identity, relative configuration, and high-level semantics. Which is exactly what a downstream action model wants.

Sigmoid Loss for Language Image Pre-Training (SigLIP)

SigLIP is a multimodal image-text model similar to CLIP. It uses separate image and text encoders to generate representations for both modalities.

SigLIP is a drop-in vision token generator in mini-VLA because our wrapper forces a clean invariant,

Any ViT vision model that can return last_hidden_state → CLS token → (B, d_model) works.

I am not going to explain this in a lot of detail, but the core idea is exactly same (for our application) as I explained in CLIP. There are subtle implementation differences that I want to highlight in the next section.

Implementing these ViT(s)

Look at this code mini-VLA/models/vision/hf_vit.py, carefully…

class HFViTVisionEncoder(VisionEncoder):
    """
    HuggingFace Vision encoder wrapper (CLIP / SigLIP).
    """
    def __init__(self, cfg: VisionEncoderCfg, hf_kind: Literal["clip", "siglip"], default_pretrained: str, mean: tuple[float, float, float], std: tuple[float, float, float],):
        super().__init__(cfg)

        try:
            if hf_kind == "clip":
                from transformers import CLIPVisionModel
                model_id = cfg.pretrained or default_pretrained
                self.backbone = CLIPVisionModel.from_pretrained(model_id)
            elif hf_kind == "siglip":
                from transformers import SiglipVisionModel
                model_id = cfg.pretrained or default_pretrained
                self.backbone = SiglipVisionModel.from_pretrained(model_id)
            else:
                raise ValueError(f"Unsupported hf_kind={hf_kind}")
        except ImportError as e:
            raise ImportError("Install transformers: pip install transformers") from e

        hidden = int(self.backbone.config.hidden_size)
        out_dim = int(cfg.d_model)
        self.proj = nn.Identity() if out_dim == hidden else nn.Linear(hidden, out_dim)

        self.register_buffer("mean", torch.tensor(mean, dtype=torch.float32).view(1, 3, 1, 1), persistent=False)
        self.register_buffer("std", torch.tensor(std, dtype=torch.float32).view(1, 3, 1, 1), persistent=False)

        self.image_size = int(cfg.image_size or getattr(self.backbone.config, "image_size", 224))

        if not cfg.trainable:
            for p in self.backbone.parameters():
                p.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.dtype == torch.uint8: x = x.float() / 255.0
        else: x = x.float()

        if x.shape[-2:] != (self.image_size, self.image_size):
            x = F.interpolate(x, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False)

        x = (x - self.mean) / self.std
        out = self.backbone(pixel_values=x, return_dict=True).last_hidden_state  # (B, N, hidden)
        cls = out[:, 0]
        return self.proj(cls)

Take a moment to absorb it as much as you can. You can find a well-documented/commented code on github. The forward pass is doing four important things:

1. Accepts mixed input types, uint8 gets scaled to [0,1] float, whereas anything else is forced to become float
2. Resizes the model’s native resolution using cfg.image_size or backbone.config.image_size
3. Normalized inside the encoder using buffers
4. Extracts a single vector token that takes in any arbitrary hidden size and returns it with fusion embedding’s acceptable size

I hope this makes sense. We are doing exactly what we did with TinyCNN in previous blogs, but this time we are doing it in a more meaningful way.

The most important difference is, in the case of TinyCNN we were training that encoder, but in the case of CLIP or SigLIP we keep their encoder weights frozen, and just use the forward pass!

So WHAT?

At this point, the integration story should feel almost suspiciously simple: every vision encoder in mini-VLA is just a token generator. In other words, the rest of mini-VLA doesn’t know that you changed the vision system — it still receives the same (B, d_model) vision token and does the same fusion + diffusion action generation. That is the entire design win.

Frozen vs trainable weights: what I actually recommend

Even though I mentioned earlier that CLIP and SigLIP weights are frozen by default, that doesn’t mean you cannot update them.

Currently we do cfg.trainable = False, so CLIP/SigLIP weights don’t get gradients. This is the cleanest baseline, fastest to train, and usually surprisingly strong.

You may set cfg.trainable = True if you have enough data and you’re confident your distribution isn’t wildly different from pretraining. Otherwise, you risk overfitting fast and losing the generality that makes CLIP/SigLIP useful.

My general observation is:

• Small robotics dataset → keep frozen
• Large, diverse dataset + strong regularization → consider finetuning

Epilogue

If you zoom out, this entire blog is really about one idea,

when you design clean interfaces, you can upgrade the internals without rewriting the system.

mini-VLA stays stable because it treats each modality as a single vector token. Today it’s CNN → CLIP/SigLIP. Tomorrow it could be DINOv2, SAM-style pooled features, or a video encoder.

So, remember this… Sometimes you don’t need a bigger framework — you need better boundaries.

If you have not read my previous blog — Building VLA model from scratch: I, which covers a philosophical overview of the same. However, this Medium-blog covers more in-depth technical portion of the same architecture, and some not-so-far-in-future scope I am working on to ‘perhaps’ generalize this model further — we will get back to that soon! Anyways, you should expect some amount of math, intuition, deeper conceptual grounding, and practical design trade-offs in this post.

This follow-up is the technical deep dive I avoided last time. We’ll go under the hood of the exact same pipeline, but this time with the math and the “why”.

[EXTRA] Why did I switch from Substack to Medium? I wanted to reference code-blocks alongside the concept and math involved. Turns out Substack doesn’t have a direct way of embedding gists into the blogs, and I wish I knew that sooner. Hence, here we are now.

Finally, let’s dive into it. If you have not read Part-1, check that out (takes about 8 mins).

Building Vision-Language-Action Model from scratch

Recap: The mini-VLA Pipeline

Vision → Language → Action pipeline

A vision encoder extracts visual features from image observations into embeddings; a vision-language projector maps these embeddings into the LLM’s representation space; a language model backbone processes the fused tokens alongside text instructions; and an action head converts LLM outputs to executable robot commands.

In other words, VLA models fine-tune VLMs for robot control, viz instead of generating captions for a given image, we generate actions for the given image and text instruction.

GitHub - keivalya/mini-vla: a minimal, beginner-friendly VLA to show how robot policies can fuse images, text, and states to generate actions

Encoding Vision, Language, and State

All encoders are defined in the mini-VLA/models/encoders.py file, and I highly recommend you to follow along with me while I explain them line-by-line.

Image Encoder

I have used a simple Convolutional Neural Network (CNN) for simplicity. If you are unaware of what that is or need a refresher, use this tutorial — PoloClub CNN Explainer.

ImageEncoderTinyCNN maps an RGB image to a single fixed-size embedding. This embedding becomes the vision token that later fuses with language/state tokens.

How does this embedding takes place? Image passes through three convolutional layers, and at each layer the spatial dimensions of the image reduce by half, and number of channels increase as 3 → 32 → 64 → 128. By the end of CNN layers, the network has built 128 learned feature maps at lower resolution — typical CNN behavior, right? Fewer pixels, more semantic channels.

Then we do a Global Average Pooling (GAP) to collapse spatial features into a 1-dimensional vector. each channel becomes a single summary statistic (“how present is feature c anywhere in the image?”). This makes the embedding translation-tolerant (like, small shifts are trivial) and of fixed-size.

At the end, we perform a Linear Projection into the model’s embedding space. This step is crucial in VLAs as it aligns vision features into the same dimensionality as other modalities (like text, or state), so these can be fused (we will come to it soon).

Despite being effective for our usecase, it is not scalable. When we do GAP, we lose grounding. Hence, modern VLAs retain multiple spatial tokens (ViT, feature maps, etc) so actions can attend to where things are, not just what is present. I am currently working on mini-VLA-v2 on the /vision branch of the same codebase to address this issue. Later, we are interested in adding cross-modal attention, and leveraging pretrained semantic vision models (like ResNet etc).

Instruction/Text Encoder

TextEncoderTinyGRU maps a tokenized instruction like “push the object to the goal” into a single fixed-size language embedding that can be fused with vision + robot state.

Well, before encoding the text, we tokenize the text using a tokenizer — mini-VLA/models/tokenizer.py, then that tokenized sentence (aka word IDs) are fed into the text encoder.

The process goes like, token IDs → X → h → z

Each token w is mapped to a learned vector in embedding space. This embedding matrix learns a continuous representation of words so that similar words (e.g., “push” vs “move”) can lie closer in embedding space. Then we use Gated Recurring Unit (GRU) to process the input sequence one token at a time, maintaining a hidden state. This is like sequence-to-vector encoder, as it compresses the entire instruction into a single vector. Lastly we do a LayerNorm to normalize each sample’s embedding across feature dimension.

This might look complicated if you don’t have good understanding of RNNs. We only use the last hidden layer, however modern VLAs keep all token embeddings and use cross-attention so actions can attend to specific words/phrases.

Robot-State Encoder

StateEncoderMLP converts a low-dimensional robot state vector (joint angles, gripper position, etc) into a fixed-size embedding that can be fused with vision and language.

This embedding acts like a state token that provides body’s innate state and kinematic context for action generation.

This module basically performs Linear(state_dim → 64) → ReLU → Linear(64 → d_model) → LayerNorm transformation.

The first Linear Layer mixes the raw state dimensions into 64 learned features. If state contains different units (radians, meters, velocities), the MLP can learn useful combinations. We all know ReLU gives non-linearity, and this matters because robot dynamics and contact regimes are highly nonlinear (remember this in future too). Then the second Linear Layer transforms current feature space into the same feature space size as your vision/language tokens, which is critical for fusion modules. LayerNorm makes the fusion stage far more stable by ensuring the state token has predictable scale.

As we push toward multi-task, long-horizon, and real-world generalization, this encoder starts to become a bottleneck. We can encode state using structured representations e.g, kinematic-aware features, graph networks, SE(3)-equivariant components, or explicit pose embeddings. Also it’s common to represent state as multiple tokens (e.g., per joint group, per limb, per sensor), enabling attention mechanisms to query what matters for the current instruction. You get the point, right? That there’s HUGE scope for improvement here as well.

Fusion as a Learned Conditioning Variable

So far we have seen a VLA policy has three modalities: vision token v, text token t, and state token s. This module fuses them into one single context embedding that becomes the condition for action head (this is important).

Refer code mini-VLA/models/fusion.py. This is the simplest possible fusion — put all modalities side-by-side so a downstream function can mix dimensions across them. And what is that downstream function? Yeah, you guessed it right. An MLP.

Now, this is pretty interesting because this fused context is the single conditioning vector c that a policy uses, eg

• diffusion head predicts noise ϵ (a, t | c)
• predicts mean action μ (c)
• autoregression that predicts action tokens conditioned on c

So, this fusion module decides what information is preserved and how it’s blended across modalities. Again, this is not the best way of fusing information because we lose spatial grounding. A better way could to use cross-attention block instead.

Diffusion Models in Action Space

I implement a conditional DDPM-style diffusion policy head for continuous actions. What does that mean? Given a fused VLA context vector c​, it learns a stochastic policy over actions a​ by modeling action generation as iterative denoising. Refer code mini-VLA/models/diffusion_head.py

Denoising model is the core network that predicts the noise, and is given by

implemented as an MLP. (this is the standard parameterization used for DDPM; the model learns to predict the noise added at time step t)

Well, Diffusion is the part which requires its own blog (the math proving it is just beautiful!), but I choose not to spoil it here. I have implemented the code almost correctly — with ONE simplified approximation which I’d like to mention (just so you don’t sue me unnecessarily)

The reverse update that I use is a simplified approximation,

instead of the full DDPM posterior mean:

However, in practice, many diffusion-policy implementations use DDIM sampling, some other scheduler, or even better posterior parameterization.

VLA Diffusion Policy

To summarize, mini-VLA is a compact vision-language-action policy that outputs continuous actions conditioned on an image, a language instruction, and the robot state vector. Refer code mini-VLA/models/vla_diffusion_policy.py. From this definition, we can say that our Objective Function is to learn a conditional action distribution

Instead of modeling this distribution directly, the policy follows the denoising diffusion probabilistic model (DDPM) formulation and learns to generate actions by iterative denoising. We have read about it eariler in this blog, and here we are just putting all these pieces together. Nothing new here.

Our Training Objective is to minimize the expected mean-squared error between the predicted and true noise

Now, the ultimate question is… Why are we predicting noise (instead of directly predicting actions)? I mean, wouldn’t that be simpler? Turns out, there are several advantages of predicting noise over actions! The noise distribution is stationary and Gaussian, learning problem is well-conditioned across timesteps, and model naturally learns a smooth conditional action manifold. This naturally leads to smoother control outputs, robustness to small perturbations and better handling of multimodal demonstrations.

CONGRATULATIONS! You just made your own VLA, from scratch!🥳

That is it, for now

Now you know just enough about mini-VLA to get started with. There is always scope for improvement, but the challenge here is to build something that works well even on the low-end devices.

Finally, it was exciting, yet exhausting to write in a new tool — it was fun mostly. I am going to share updates, like implementing ViT and all in the updates SOON, so stay tuned! I hope you found this both useful and enjoyable. Your feedback on the previous blog helped shape this one, and I’m grateful to everyone who shared thoughtful, constructive criticism.

I love talking about Robotics and AI, and writing these blogs helps me connect with like-minded people. Thank YOU for taking the time to learn and grow today.

:)