DEV Community: Glen Yu

What I learned training an ACT policy for the LeRobot SO-ARM101

Glen Yu — Wed, 02 Sep 2026 13:54:52 +0000

I’ve been learning a lot about robotics in the last little while, so I’m here to document my journey and share my learnings starting with Action Chunking with Transformer (ACT) policies.

Introduction

I have been fascinated by robotics for a long time, but for years the space felt fragmented. There were plenty of robot arms on the market, but rarely a cohesive framework or active open-source community tying hardware and modern machine learning together. That changed when I discovered the SO-ARM101 within Hugging Face’s LeRobot ecosystem. A low-cost, 3D-printed arm backed by standardized datasets, accessible imitation learning pipelines, and an active community gave me the confidence to finally take the leap into physical AI.

And what a journey it has been. From initial setup quirks to training my first Action Chunking with Transformer (ACT) policy, getting everything to run reliably taught me plenty of lessons that are not spelled out in the official documentation. Here is everything I learned along the way so you can skip the trial and error.

Setup

Development environment

I use pyenv on my MacBook and set up a miniconda3-3.12 environment as the base. The initial setup steps outlined in the official documentation were accurate and straightforward. I installed the LeRobot package from PyPI with the necessary extras for CLI scripts and Feetech bus servo control:

pip install 'lerobot[core_scripts,feetech]'

How many cameras do you need?

When designing your physical workspace, I recommend at least two cameras: a wrist camera and an overhead or side camera.

Wrist camera: Provides fine-grained visual feedback on where target objects are relative to the gripper. Much like a racehorse with blinders, its field of view is narrow and focused entirely on the manipulation point.
Top/Overhead camera (Bird’s Eye view): Gives the model global context for where the arm and objects sit on the table. An overhead perspective avoids the blind spots created by side-mounted cameras, where the arm body can obstruct the view.
Front-facing base camera (Optional third): Depending on your task complexity, adding a camera at the base facing forward provides depth perspective that complements the top-down and wrist views.

I have seen basic budget setups use the stock 15 FPS wrist camera paired with a built-in MacBook webcam, but a dedicated multi-camera layout yields far better policy consistency.

Hardware Selection

2 x InnoMaker 1080P USB2.0 UVC camera: A drop-in upgrade for the stock 15 FPS wrist camera that came with the arm. Its compact form factor and optional aluminum case also made it easy to mount at the arm base as a front-facing camera.
Pro HD 1080P/30FPS webcam with tabletop tripod: Used for the bird’s eye view. Placing it on a nearby shelf with the tripod made it easy to dial in the height and tilt angle without buying specialized clamp mounts.

Dataset recording and thread tuning

When recording episodes, allocate at least two encoder threads per camera (dataset.encoder_threads=2) to avoid dropping frames during teleoperation. For my 3-camera setup, I assigned 3 threads per camera, utilizing 9 CPU cores on my MacBook M3 Pro.

While running cameras at 15 FPS reduces dataset size and CPU overhead, modern USB cameras natively default to 30 FPS. Standardizing on 30 FPS across all sensors avoids framerate mismatch issues during episode logging.

Downscaling resolution for faster training

While 1080p cameras provide crisp video for human eyes, high-resolution feeds create bloated datasets, spike VRAM usage, and slow training to a crawl. Standard computer vision backbones in ACT do not need full HD detail to learn spatial features. Resize your camera feeds to lower resolutions like 640x480 for both training and policy evaluation. This will keeping storage requirements low and training fast.

Training strategy

For your first ACT policy, eliminate as much physical and visual variability as possible. Use a single-color Lego brick (or build a slightly larger block using bricks of the same color) and place it against a clean, uniform background. For the target container, choose a simple box or bowl with a matte, single color. Highly reflective, translucent, or patterned objects add unnecessary visual noise that makes early policy convergence much harder.

NOTE: Before settling on a Lego brick, I initially tested with a cube-shaped USB wall charger. This turned out to be a poor choice. Depending on how it rested on the table, the prongs and varying face textures created significant visual variability from different angles. That extra visual noise impacted the gripper’s consistency in gripping the block.

How much training is needed?

A common question when moving from classical deep learning to imitation learning is how training steps map to traditional epochs. In a 3-camera setup, 100 recorded episodes typically yield around 45,000 rows of Parquet data containing synchronized camera images and servo motor joint states (each row of data corresponds to a single frame of video). If you train for 20,000 steps with a batch size of 32, the relationship is:

Epochs = Total Steps x Batch Size / Dataset Size

Plugging in those numbers:

Epochs = 20000 x 32 / 45000 = 14.2

An epoch represents one full pass through your training data. In standard PyTorch or TensorFlow workflows, 14 epochs usually still falls a bit short for a vision-based policy to generalize; 25 to 50 epochs is a more reliable baseline for basic convergence.

Task complexity vs. dataset size

Epoch count is only one metric. Training for hundreds of epochs will not improve policy quality if your underlying dataset lacks sufficient coverage. You want a dataset with broad enough coverage to match the complexity of your task, and then pair it with enough training steps to allow the model to converge:

Fixed start and fixed target (Low variance): If the arm picks up a block from a static Point A and drops it into a container at a static Point B, a modest dataset of 30 episodes trained for 15 epochs can yield a functional baseline policy.
Randomized start and fixed target (Medium variance): If the Lego block can appear anywhere within the workspace while the container stays fixed, you will likely need around 100 episodes trained for 30 epochs to properly cover the workspace distribution.
Randomized start and randomized target (High variance): If both the block and the container change positions across every demonstration, you should aim for 250+ episodes trained for ~50+ epochs so the Transformer encoder can learn generalized spatial relationships rather than memorized trajectory paths.

Interpreting L1 Loss

When monitoring your training curves, ‭L1 Loss serves as your primary convergence indicator. While there is no universal cutoff for success, ACT policies for manipulation tasks typically begin showing reliable trajectories when the ‭L1 Loss falls between 0.01 and 0.05.

What does an ‭L1 Loss of 0.02 actually mean in practice? It represents the mean absolute error between the model’s predicted trajectory chunk and the ground-truth teleoperated actions across normalized joint space. Keep in mind that ‭L1 Loss is an aggregated average across all degrees of freedom. On a 6-DOF arm, five joints might track with near-zero error while a single critical joint (such as the wrist pitch or gripper closing action) exhibits higher error, yet the overall average can still look deceptively low at 0.02. Ultimately, running real-world policy rollouts will always be the true test of performance, and seeing your arm execute the task autonomously is easily the most rewarding part.

My approach

My goal was to train an ACT policy capable of picking up a Lego block and placing it inside a bowl, where both the block and the bowl could be placed at arbitrary positions across the workspace. I set a target success rate of 75% or higher.

I recorded all demonstration episodes locally on a makeshift workbench. Because the LeRobot pipeline integrates seamlessly with the Hugging Face Hub, I pushed the datasets upstream and ran the resource-heavy training jobs on an NVIDIA DGX Spark at the office. This allowed me to train with larger batch sizes and higher step counts far faster than running on my MacBook.

I started with a medium-variance baseline: 50 episodes (~21.5k frames) where the Lego block was placed in different locations around the workspace while the destination bowl remained fixed in one spot.

Training configuration: 20,000 steps, batch size of 32 (~30 epochs).
Compute: Training took roughly 8 hours. A batch size of 32 consumed approximately 22GB of VRAM.
Result: (50 episodes for 20k training steps, L1 Loss = 0.074)

With the single-target baseline proven, I introduced full spatial variance. I recorded 150 new episodes where the positions of both the Lego block and the bowl changed with every demonstration, then merged them with the initial dataset for a total of 200 episodes (~83.5k frames).

Training configuration: 100,000 steps, batch size of 32 (~38 epochs).
Compute: Training time was roughly 41 hours.
Result: (200 episodes for 100k training steps, L1 Loss = 0.052)

While the v1 model performed well overall, I wanted to push reliability higher. I recorded an additional 125 episodes specifically covering awkward placements, tight table boundaries, and difficult angles. Merging these gave a final dataset of 325 episodes (~141k frames).

Training configuration: 250,000 steps, batch size of 32 (~57 epochs).
Compute: Total training time ran for approximately 104 hours.
Result: (325 episodes for 250k training steps, L1 Loss = 0.040)

Training for 104 hours might sound like overkill, but leveraging the DGX Spark compute allowed the Transformer to generalize smoothly across tricky spatial boundaries and hit the 75%+ success target consistently.

What’s next?

Working through the “hello world” of robotic manipulation with the SO-ARM101 was a lot of fun, but it is just the starting point. Next, I plan to tackle more complex, deformable object manipulation tasks, such as folding clothes with an ACT policy.

Beyond pure imitation learning, I am also excited to experiment with Vision-Language-Action (VLA) models such as Hugging Face’s SmolVLA, Physical Intelligence’s Pi, and NVIDIA’s GR00T to see how generalist multimodal backbones handle zero-shot spatial reasoning and natural language task prompting.

(And of course, I plan to share my learnings with these tools as well)

What I would love to see

From a tooling perspective, native macOS support for simulation environments like NVIDIA Isaac Sim would be a huge quality-of-life improvement for Apple Silicon developers. Currently, Isaac Sim remains restricted to Linux and Windows with dedicated NVIDIA RTX hardware.

In an ideal setup, I would love to have a dedicated local AI workstation at home, powered by an NVIDIA DGX Spark or an AMD Ryzen AI Halo. Having native Linux and serious local compute would let me run GPU-accelerated physics simulations in Isaac Sim and train large physical AI policies right from my home workbench.

If you found this post helpful or are currently troubleshooting your own LeRobot build, feel free to connect or share your setup in the comments below.

Google Colab, but in your favourite terminal

Glen Yu — Fri, 05 Jun 2026 21:45:00 +0000

While some of my recent posts have involved using the Colab extension for VS Code and the Antigravity IDE, I actually prefer working in the terminal and Vim. The new Colab CLI finally lets me work in my natural habitat, and it opens the door for autonomous workflows!

Setup

Currently, installation is handled via pip or uv. It's straightforward, though, I'm holding out hope for a brew formula in the future:
uv tool install google-colab-cli

I'm testing Version: 0.6.dev7+g510115b0c inside Ghostty. The Colab CLI is pretty solid, but I do have some feedback and nitpicks I'd like to share (but more on that later).

Creating a new session

Creating a session is simple: colab new [-s SESSION_NAME] [--gpu T4|L4|A100|H100] [--tpu v5e1|v6e1]:

SESSION_NAME: This is optional. If you leave it blank, the CLI generates a random unique ID for you.
--gpu and --tpu: The hardware accelerator flags are optional, but omitting them defaults to a standard CPU-only instance. The specific accelerator chips you can request depend on your Colab tier, which you can check via colab pay.

NOTE: If you only have one active session, the CLI targets it by default. This makes the -s flag unnecessary for subsequent commands.

Testing Colab CLI's capabilities

CLI certainly sounds cool, but how does it handle artifacts and images? More importantly, how debuggable is it? I decided to find out by running a Fashion MNIST PyTorch example.

Handling artifacts

To get started, I installed my requirements using colab install torch torchvision matplotlib. If you prefer a more standard approach, you can also use colab install -r requirements.txt.

Once the environment was ready, I executed the training script using colab exec -f ./fashion_mnist_TRAIN.py and here's the output:

[colab] Using unique session '8c860c'.
Using CUDA device.

Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28])
Shape of y: torch.Size([64]) torch.int64
NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)
Epoch 1
----------------------------------
loss: 2.299284 [   64/60000]
loss: 2.296870 [ 6464/60000]
loss: 2.278877 [12864/60000]
loss: 2.274183 [19264/60000]
loss: 2.266371 [25664/60000]
loss: 2.233031 [32064/60000]
loss: 2.227577 [38464/60000]
loss: 2.201082 [44864/60000]
loss: 2.191893 [51264/60000]
loss: 2.167733 [57664/60000]
Test Error:
 Accuracy: 52.6%, Avg loss: 2.163899

...
...

Epoch 25
----------------------------------
loss: 0.495004 [   64/60000]
loss: 0.619403 [ 6464/60000]
loss: 0.404403 [12864/60000]
loss: 0.658543 [19264/60000]
loss: 0.587690 [25664/60000]
loss: 0.565604 [32064/60000]
loss: 0.592686 [38464/60000]
loss: 0.675307 [44864/60000]
loss: 0.645753 [51264/60000]
loss: 0.571738 [57664/60000]
Test Error:
 Accuracy: 79.9%, Avg loss: 0.574529

Done!
Saved PyTorch Model State to 01_mnist_model.pth

When the training finished, the PyTorch model is saved in the Colab instance. To manage these files, you have two solid options:

Google Drive: Use colab drivemount [PATH] to mount your Google Drive to save artifacts directly to a persistent cloud filesystem.
Local Download: Use colab download REMOTE_FILE LOCAL_FILE to pull the model back to your local machine.

Working with images

I was pleasantly surprised to find that images displayed just fine! Despite being a self-proclaimed shell aficionado, I had rarely worked with images in a terminal environment and didn't realize how seamless this could be.

When I ran the inference script with colab exec -f ./fashion_mnist_PREDICT.py, and the CLI automatically handled the visual output:

NOTE: Here are the links to the training script and inference script.

Debugging

Debugging Colab jobs is arguably easier than before because you can jump directly into a running instance using colab console:

Once you shell into the environment, you can inspect downloaded files, verify artifacts, or install additional dependencies on the fly.

If you need to keep a record of your work, you can export your logs and session history using colab log -o SESSION_HISTORY_FILE. This is particularly useful if you want to save your history directly into a mounted Google Drive filesystem for long-term storage or later review.

COLAB_SKILL.md

The resurgence of CLI tools is driven by their efficiency over traditional GUIs. While GUIs are designed for human intuition and rely on visual cues, these elements represent noise that an AI agent must filter out.

LLMs are already great with text and hence the input & output formats of a CLI are ideal for agents and machines to interpret. In fact, the most powerful way to leverage the Colab CLI is through autonomous agents! Colab CLI's repo provides a COLAB_SKILL.md file for you to get started for that agentic integration.

Issues & nitpicks

As previously mentioned, the Colab CLI is still being developed, so there's bound some issues here and there. Here are some that I've noticed and would love to see fixed:

Session Timeouts: Relatively short timeouts and automatic pruning can disrupt productivity. While my personal impact has been minimal, this can be a significant hurdle for users with longer running tasks.
Interface Disconnects: CLI sessions shows up as "Unknown notebook" in Colab Web and are invisible in the Antigravity IDE or VS Code extension. I've filed a GitHub Issue detailing this lack of cross-tool visibility.
Hardware Parity: Colab Web and IDE extension now offers NVIDIA G4 (Blackwell) instances with 96GB of VRAM. To support high-performance AI workloads, the CLI should offer parity with these newer hardware tiers.

Summary

The Colab CLI is a significant step forward for developers who prefer the speed and focus of the terminal over a traditional browser-based UI or IDE. By bringing the power of Google's cloud hardware into a local, shell-based environment, it successfully bridges the gap between high-performance computing and a minimalist local workflow.

While the tool is clearly still in its early stages and has a few rough edges to smooth out, the core functionality is already impressive. The ability to manage remote instances, handle graphical output, and shell into a VM as easily as a local process makes it a formidable addition to any developer's toolkit.

Whether you are looking to automate your training pipelines or simply want to stay inside Vim while you scale your models, the Colab CLI offers a glimpse into a more efficient, autonomous future.

ML acceleration guide: TPUs vs GPUs

Glen Yu — Tue, 28 Apr 2026 00:16:10 +0000

There’s a lot of hype around GPUs and NVIDIA, but how much do you know about TPUs?

Article includes code examples you can find near the end

Rise of GPUs

Graphics Processing Units have been around for quite some time and their job is to render 2D and 3D graphics in to millions of pixels, calculating their colour, texture, lighting, in parallel to send to your monitor. For a 60Hz monitor that means producing rendered frames 60 times every second.

Rendering graphics is one thing, but developing code for handling GPUs was a little more difficult. That is, until NVIDIA launched CUDA (Compute Unified Device Architecture) in 2006, which allowed scientific researchers and developers who work in fields that require massive parallel math to take advantage of a GPU’s capabilities. With the rise of machine learning in the early 2010’s, it was discovered that the massive parallel math was exactly what ML engineers needed to train deep neural networks. Since then, the focus of CUDA has been shifting more towards optimizing for machine learning and AI workloads.

Because GPUs were commercially available and relatively inexpensive at the time, the barrier to entry was low. An ML engineer could train models on their NVIDIA graphics card during the day and jump into a game of League of Legends at night on the same hardware.

Honourable mention

AMD’s GPUs with Radeon Open Compute (ROCm) in an open-source software stack designed to compete in the AI ecosystem. Though it’s not as popular as CUDA, this gap is closing with Meta recently signing a deal to expand its existing partnership with AMD.

Tensor Processing Unit

In the early 2010s, Google projected that the growing demands of its AI workloads, particularly the rapid adoption of deep learning across products like Search and Photos, would require doubling its data center computing capacity roughly every year and a half. Rather than scale generic hardware indefinitely, Google sought a more efficient solution purpose-built for neural network computation, and thus the Tensor Processing Unit (TPU) was born. The TPU is a custom application-specific integrated circuit (ASIC) designed by Google specifically to accelerate AI workloads, deployed internally starting in 2015. By specializing the hardware for the dense matrix operations at the heart of neural networks, TPUs achieve dramatically better performance per watt than general-purpose CPUs or GPUs, reducing both energy consumption and cooling demands at data center scale.

Google has a tradition of making tools it uses internally available to the broader world, and TPUs are another example of this. The existence of TPUs was first publicly announced at Google I/O in 2016. In 2018, Cloud TPU v2 became available for external users through Google Cloud, marking the first time developers outside Google could harness the same accelerators powering Google’s own AI systems. TPUs also come in two performance flavours: efficiency and performance to meet different market needs.

NOTE: As of the 8th generation of TPUs announced during Google Next 2026, efficiency and performance TPUs will be renamed inference and training respectively in favour of a more descriptive, workload-based naming convention.

Architecture layout

From an architectural standpoint, GPUs can be thought of as being individual computers with accelerators (picture your home gaming PC). If you want to connect them into a cluster, it would be over network, but no matter how fast the network is, it still has to cross node boundaries, and bandwidth drops as a result.

TPUs are designed from the ground up to be interconnected at a massive scale with a physical layout that involves thousands of TPU chips in a torus topology which gives every chip 6 neighbours (two per axis, one on each side). Recognize that interconnect bandwidth would be the main bottleneck at this scale, Google designed their own proprietary Inter-Chip Interconnect (ICI) network which provides uniform, high-bandwidth, low-latency connections between all the chips in a slice regardless of physical location. With torus topology, there is no concept of crossing a node boundary. When you request TPUs, you do not get the entire TPU cluster or pod. Rather, you get only a small subset or slice. To make this possible, Google developed Optical Circuit Switch (OCS) to be able to rewire physical connections on the fly (entirely in software), allowing the same hardware to serve different workload shapes without any physical reconfiguration.

NOTE: Efficiency TPU versions use a 2D torus topology, while Performance TPUs leverage a 3D torus architecture to give you maximum performance with minimum latency.

Precision and range

A floating-point number consists of three parts:

Sign: Positive or negative (represented by the first bit)
Exponent: Determines the range of the number
Mantissa: Significant digits of a floating-point number, which determines the accuracy

Traditionally, the standard for high-performance computing was FP32. When AI researchers moved to FP16 to save memory, they lost more than just accuracy: they also lost range. FP32 uses 8 bits for the exponent, while FP16 uses only 5. The 3-bit difference in the exponent bits amount to an almost 10³⁴ difference in range (FP32 has a range of 3.4 x 10³⁸, while FP16 only has a range of 6.5 x 10⁴). In deep learning, where gradients can be incredibly tiny, FP16 often suffers from underflow (meaning numbers are being rounded to 0 because it is too small for FP16’s range to represent), requiring a technical workaround called “loss scaling” to keep the math stable.

Google Brain (now part of Google DeepMind) solved this invented Brain Floating Point (bfloat16), which simply shifted 3 bits from the mantissa to exponents:

Format	Total Bits	Exponent Bits	Mantissa Bits
FP32	32	8	23
FP16	16	5	10
bfloat16	16	8	7

By sacrificing precision for range, bfloat offers the same massive range as FP32, but with the reduced memory and bandwidth of FP16. A huge reason for why this works is that deep learning models are surprisingly noise-tolerant and having more training stability is for more important than having a few extra decimal places of precision. Today, bfloat16 is the de facto standard for training modern LLMs on NVIDIA’s GPUs and Google’s TPUs.

Why XLA matters

Standard Python execution typically takes an eager approach. This means it executes each step as it is being encountered. This is great for debugging because you can insert print statements to inspect variables at any point.

XLA (Accelerated Linear Algebra), on the other hand, is a domain-specific JIT compiler. Instead of executing steps one by one, it analyzes the entire execution graph to optimize and fuse operations before they run. This lazy approach creates an initial warm-up delay, but once the training starts, it is significantly faster than standard methods. The tradeoff is transparency: your step-by-step Python code becomes an optimized “black box”, making traditional debugging strategies more difficult. This is why TPUs are powerhouses for massive enterprise training, while GPUs remain the flexible choice for quick experimentation.

NOTE: Though XLA was built for TPUs, it’s also made its way into the NVIIA GPU ecosystem via tools such as JAX and torch.compile since PyTorch 2.0.

TorchTPU

Google is engineering a TorchTPU stack that will provide native PyTorch support. This would allow you to run models in TPUs as they are with full support for native PyTorch features. TorchTPU is currently in preview, and once it becomes GA, you can be sure I’ll be diving deeper into it!

Code example

I’m including a couple of Jupyter notebooks that I ran via Antigravity + Colab plugin for you to try yourself:

As you will see from the results below, TPU is indeed faster. However, my example isn’t large enough or complex enough to really showcase the true speeds that TPU can bring.

NOTE: I have a Colab Pro account which affords me access to additional GPUs and TPUs. The Colab free tier only includes NVIDIA T4 and TPU v5e-1

Interpreting training results

These are some benchmark trainings (epochs: 50, batch size: 512) in which I used a NVIDIA T4 GPU with (default) FP32 vs Google TPU v5e-1 (single chip TPU) with bfloat16. As expected, TPUs were faster but with lower precision:

I then trained the same model using the T4 GPU using bfloat16 but noticed a massive performance drop. This was due to the T4 being an older generation GPU that did not support bfloat16 natively and had to emulate which added a lot of overhead. Switching to a newer L4 GPU, I was able to see the (tiny) performance gain along with the reduced precision:

Finally, I thought I’d see how the training would perform on a newer TPU v6e-1 and I was blown away by the improvement:

Conclusion

Comparing GPUs and TPUs isn’t exactly apples-to-apples. They represent fundamentally different philosophies in architecture, memory management, and execution.

In the modern enterprise, it isn’t usually a matter of choosing one over the other, but rather using each where it shines. For rapid iteration and smaller workloads, the flexibility of GPUs is unmatched. However, once a project hits a certain scale, the domain-specific architecture of the TPU becomes the clear winner in efficiency and throughput.

TPUs are as fast as they are because they are a specialized one-trick pony, but to truly harness that power requires a deeper understanding of the stack. The biggest challenge isn’t often the compute itself, but rather: “How do I feed data to the TPUs fast enough and efficiently enough so that it doesn’t become the bottleneck?” and ensuring your input pipeline is fast enough so that the hardware doesn’t sit idle.

In future posts, I'll dive deeper into these advanced concepts to show how you can optimizing data pipelines to get the most out of your TPUs.

BONUS: Google’s 8th-generation TPUs announced at Google Next

Implementing a RAG system: Run

Glen Yu — Tue, 07 Apr 2026 17:00:00 +0000

In the "Crawl" and "Walk" phases, I introduced the basics of RAG and explored ways to optimize the pipeline to increase efficiency and accuracy. Armed with this knowledge, it's time to productionize our learnings.

Run

In the "Crawl" and "Walk" phases, we explored RAG fundamentals using local tools, proving how much document processing and re-ranking impact performance. While you could certainly scale those manual workflows into production, do you really wan to manage the infrastructure, data pipelines and scaling hurdles yourself?

Welcome to the "Run" phase. Here we leverage Google Cloud's Vertex AI RAG Engine - a fully managed solution that automates the entire pipeline so you can focus on building, not maintenance.

Vertex AI RAG Engine

Vertex AI RAG Engine is a low-code, fully managed solution for building AI applications on private data. It handles the ingestion, document processing, embedding, retrieval, ranking, and grounding to ensure that the response is highly accurate and relevant.

Optional document pre-processing with Docling

Though RAG Engine comes with its own parser options, I still opted to pre-process my documents using Docling first and upload them to a Google Cloud Storage bucket for ingestion:

if __name__ == "__main__":
    DATA_URL = "https://storage.googleapis.com/public-file-server/genai-downloads/bc_hr_policies.tgz"
    DATA_DIR = download_and_extract(DATA_URL)

    docling_docs = Path("./docling_docs")
    docling_docs.mkdir(parents=True, exist_ok=True)

    print("> Processing PDFs...")
    data_path = Path(DATA_DIR)
    pdf_files = list(data_path.glob("*.pdf"))

    for file in pdf_files:
        try:
            result = converter.convert(file)
            markdown_content = result.document.export_to_markdown()

            docling_doc_name = Path(file).stem + ".md"
            with open(f"docling_docs/{docling_doc_name}", "w", encoding="utf-8") as f:
                f.write(markdown_content)
        except Exception as e:
            print(f"Error on {file}: {e}")

    print("> Uploading Docling docs to GCS...")
    upload_folder_to_gcs("./docling_docs", GCS_BUCKET, GCS_BUCKET_PATH)

One benefit of doing it this way is I get more visibility and control over the document processing step and can validate the contents from the original PDFs with the Docling documents (Markdown). This allows me to use the Default parsing libraries option, which is also free compared to the LLM parser and Document AI layout parser options which have an additional cost and setup component to them.

I do lose out on the benefits of the hybrid chunking strategy that Docling would have provided (as seen in the "Walk" phase), because that is determined by the layout parser that I choose here. If I wasn't using Docling, I think the LLM parser would be the parser option that I'd be gravitating towards.

Vector database options

When it comes to vector database options, you'll see several choices in the RAG Engine menu (including regional "Preview" features). I chose the RagManaged Cloud Spanner (also referred to as the RagManagedDb) because it offers the fastest path from data to insights with the least amount of infrastructure management. While Spanner is typically an enterprise-grade database, the RAG Engine allows you to spin it up on the Basic tier. This allocates 100 processing units, which is 10% of a Spanner node, making it perfect for smaller datasets while still giving you the reliability of a managed service without the enterprise-grade cost.

IMPORTANT: Even on basic tier, this will still run you about $65 USD/month, so please remember to delete and clean up this RAG corpus once you're done experimenting with it.

For those prioritizing flexibility, the RAG Engine also supports third-party options like Pinecone and Weaviate. These are excellent choices if portability is a requirement, allowing you to maintain a consistent vector store even if you decide to shift parts of your RAG stack to a different cloud provider or platform later on.

Ranking & grounding included

Once the RAG corpus is created, you can perform some manual testing to validate. When you ask RAG Engine for search results, re-ranking and grounding is done automatically to ensure relevance and correctness:

Model Armor

In a production setting (especially if it's going to be public facing), you will want guardrails. I've written about Guardrails with Agent Development Kit in the past and is implemented through callbacks within ADK. It works the same way here and can be used to inspect text as it flows into and out of the LLM/agent. Key capabilities include:

Prompt injection & jailbreak detection: Attempts to trick the AI into ignoring instructions
Sensitive Data Protection: Natively integrates with Google's Data Loss Prevention (DLP) to scale for various types of sensitive information (PII)
Malicious URL detection
Responsible AI (RAI) filters: Hate speech, harassment, dangerous, and sexually explicit content

I configured my Model Armor policy template which I then invoked via the Python SDK library to determine whether the text was safe.

Updated example

You can find the code for "Run" phase → here

Querying the vector database (RAG Engine in this case) is a lot less involved as I don't have to write a many of the logic to pass the semantic search results to a re-ranker, because RAG Engine takes care of all of that for me!

I once again ran the same two benchmark questions as I did in the "Crawl" and "Walk" phases:

I threw in a couple of extra questions to make sure Model Armor wasn't sleeping on the job, but overall I liked the detail and accuracy of the answers I was provided.

NOTE: In my updated example, I only added a before_model_callback, meaning I'm only checking the input prompts and not the response. An after_model_callback should be implemented to ensure the generated response is also scanned and preventing the AI from accidentally leaking sensitive internal data it might have retrieved from the RAG corpus (I omitted the output check here simply because I know there's no sensitive data in this particular dataset).

Summary

The purpose of this "Crawl, Walk, Run" series was to take you on a journey from managing code to delivering value. In the earlier phases, we deconstructed the mechanics of RAG works to understand the roles that chunking, embedding, re-ranking play in the overall system. In this final phase, we see how Vertex AI RAG Engine and Model Armor streamline those manual components. By offloading infrastructure management and safety logic to Google Cloud's managed services, you can ensure your system is scalable, accurate, and secure from day one.

Next Steps

Currently in private preview is Vector Search 2.0 with RAG, but reading through its documentation and features, it looks pretty interesting, so once it becomes GA, I will definitely give it a try.
I'm also looking forward to all the new AI-related announcements that is sure to happen at Google Next!

Additional learning

Interested in finding out more about how to secure your agent by sanitizing input and output? Try out this Model Armor Codelab!
Vertex AI RAG Engine isn't your only option for a managed RAG, if you'd like to try a different option that uses Vector AI Search, might I suggest the Building Agents with Retrieval-Augmented Generation Codelab?

Happy learning!

Implementing a RAG system: Walk

Glen Yu — Tue, 31 Mar 2026 17:00:00 +0000

Now that we've established the basics in our "Crawl" phase, it's time to pick up the pace. In this guid, we'll move beyond the initial setup to focus on optimizing core architectural components for better performance and accuracy.

Walk

We ended the previous "Crawl" design with a functioning AI HR agent with a RAG system. The responses, however, could be better. I've introduced some new elements to the architecture to perform better document processing and chunking, as well as re-ranker model to sort the semantic retrieval results by relevance:

The ugly Docling

IBM's Docling is an open-source document processing tool and easily one of the most effective ones I've tested. It can convert various file formats (e.g., PDF, docx, HTML) into clean, structured formats like Markdown and JSON. By integrating AI models and OCR, it doesn't just extract text, but also preserve the original layout's integrity.
Through its hierarchical and hybrid chunking methods, Docling intelligently groups content by heading, merges smaller fragments for better context, and attaches rich metadata to streamline downstream searching and citations. Here's a Python function I use for chunking a PDF file:

def docling_chunk_pdf(file: str) -> tuple[list, list]:
    converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(
                pipeline_options=pipeline_options,
            )
        }
    )

    result = converter.convert(file)
    doc = result.document
    chunker = HybridChunker()
    chunks = list(chunker.chunk(doc))
    chunk_texts = [c.text for c in chunks]

    return chunks, chunk_texts

I plan to take a deeper dive into Docling in a future article, so give me a follow so you won't miss it! 😄

Dot product vs cosine similarity

In the "Crawl" post, I talked briefly about cosine similarity and how it ignore magnitude and only focuses on the angle between two vectors. This is because normalization is baked into the cosine similarity formula.
Dot product is effectively cosine similarity but without the final normalization step, which is why its result is affected by the magnitude of the vectors. Since many modern embedding models output pre-normalized unit vectors, the extra normalization step in cosine similarity becomes a redundant calculation. By using dot product on these pre-normalized vectors, you can achieve identical results with higher computational efficiency.

NOTE #1: While switching to dot product can increase your raw retrieval throughput, the latency gains may feel negligible when considering the entire end-to-end RAG pipeline depending on your particular use case and scale.

NOTE #2: A friendly reminder that choosing dot product over cosine similarity has the hard requirement that your vectors be normalized beforehand, or the magnitude will skew your search results. It's also quite easy to update your search configuration to use one or the other. If you're ever in doubt, just run a quick test with both settings to verify that both methods return the exact same nearest neighbours (top semantic matches).

Re-ranking

Standard search is built for speed and not deep understanding, so it can sometimes miss nuances. Re-ranking takes a crucial "second look" at the standard retrieval results to see which one(s) actually address the user's query. While the Cosine distance represents how similar the query and document align in the vector space, a "close" match doesn't guarantee an answer. The re-ranker's job to is to bridge this gap by scrutinizing the top results to assign a true relevance_score and ensure the most helpful contexts rise to the top. Here's what that snippet of code looks like:

co = cohere.ClientV2()

response = co.rerank(
    model=RERANKING_MODEL,
    query=user_query,
    documents=documents_to_rerank,
    top_n=len(candidate_responses),
)

reranked_results = []
for res in response.results:
    original_data = candidate_responses[res.index]

    reranked_results.append({
        "content": original_data["content"],
        "source": original_data["source"],
        "heading": original_data["heading"],
        "page": original_data["page"],
        "search_distance": original_data["search_distance"],
        "relevance_score": res.relevance_score
    })

As part of the full code that performs the re-ranking, I assign a threshold for the relevance score. Scores lower than this threshold is deemed irrelevant.

Updated example

I'm shaking up the stack for the "Walk" phase! In addition to using a different document processor, I will also be using a different embedding model and vector database.
Since I wanted to try out Cohere's re-ranking model, I opted to lean into their full suite and use their embedding model as well. I made a deliberate choice here to set the embedding dimension to 384, which is a lower than the 768 I previous used in the "Crawl" example. I wanted to handicap the initial semantic search, and by doing so, we can more clearly see the re-ranker work its magic to fix the order or the results.
I switched out ChromaDB with LanceDB to showcase just how many robust, easy-to-use open-source local vector databases are available for use.

Querying the HR agent

While I kept the core agent configuration from the "Crawl" phase the same, the addition of the re-ranking step made a significant impact. I asked the same two benchmark questions and this time the results were more refined and accurate:

You can find the code for the "Walk" phase → here

Next steps

Now that we've manually optimized our retrieval and re-ranking, the next step is to scale. I will be migrating this architecture to Vertex AI's RAG Engine for a fully managed, high-performance RAG pipeline at an enterprise scale.

Additional learning

I used Cohere's embedding and re-ranking models in my example, but if you want to try out Vertex AI's re-ranking capabilities (and more), try out this Advanced RAG Techniques Codelab.

Implementing a RAG system: Crawl

Glen Yu — Tue, 24 Mar 2026 18:00:00 +0000

I'm starting a "Crawl, walk, run" series of posts on various topics and decided to start with Retrieval-Augmented Generation (RAG). Learn the basics and progress to a production-ready system!

Crawl

In this phase of your journey, we're going to learn about the core concepts of a Retrieval-Augmented Generation (RAG) system and then apply them in a simple example.
We're going to build a Human Resources (HR) agent that can help answer and navigate HR-related questions. Using the Government of British Columbia's HR Policy PDFs as our knowledge base, we will process, chunk, and embed the documents into a local vector database. This allows the agent to provided grounded answers and ensures that every response is rooted directly in the ingested BC government policies.

Why RAG?

RAG is a very common design pattern that turns a standard LLM into an informed AI agent. Standard models can be a "black box", but RAG gives your agent an "open-book test". It bypasses knowledge cutoffs by linking directly to your documents, providing factual grounding and citations. No fine-tuning is required, data can be updated quickly. RAG provides a real-time bridge between your LLM and your data.
This post will focus primarily on indexing and retrieval of the your data. Let's get started!

How do you eat an elephant? One bite at a time

It's not feasible to have to feed all the information into the AI every time you want to ask it a question. Instead, it is broken down into smaller, more manageable pieces called chunks, which the AI can process and retrieve efficiently.
We will use a "recursive character chunking" strategy which is a fast and smart and will try to split at natural boundaries like paragraphs, but can still cut off mid-sentence if the chunk is too big. An overlap is used to ensure that context isn't lost at the edges of a the cut if a split does occur.

Snippet of code used for splitting & chunking using LangChain:

from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

loader = DirectoryLoader(
    DATA_DIR,
    glob="./**/*.pdf",
    loader_cls=PyPDFLoader
)
docs = loader.load()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=100
)
chunks = text_splitter.split_documents(docs)

Recursive character chunking is the successor to "fixed-sized chunking", which is just a fixed sliding window. Here, you always need the overlap because you never know how much of what sentence you're cutting off.

What's the vector representation of 'Life'?

The embedding process transforms text chunks into vectors, which are mathematical arrays of floating-point numbers that capture semantic meaning. However, higher dimensionality won't necessarily generate better results. For simple or straightforward documents, expanding the vector size often introduces latency and computational overhead without providing better search accuracy.

Two sides of the same coin: Indexing & retrieval

Indexing and retrieval are two parts of the same conversation, and you must use the same embedding model for both. This is important because every embedding model puts emphasis on words in a sentence differently. One might prioritize the subject, while another might prioritize the action, which would yield different results.

What happens when different embedding models try to embed “To be, or not to be…”:

Selection the appropriate embedding type is also very important. The documents that you embed and index are usually a long, structured documents where the focus is on the information it provides. This is in contrast to the user queries which are usually short, messy text, so the retrieval process focuses on the information it is looking for.

Finding a match

Once your user query is embedded, the RAG system performs a similarity search against the vector database to identify the most relevant answers. In most vector databases, this is calculated using cosine similarity. This metric focuses exclusively on the angle between vectors rather than their magnitude; it measures how closely the semantic "intent" (angle) of the query aligns with the document, regardless of the text's length or word frequency. This is important because it means the AI can recognize that a short question and a long technical manual can share the same intent even if their scale (magnitude) is completely different.

Putting it all together

Link to my GitHub repository → here

To handle HR questions, I'm building an agent using Google's Agent Development Kit (ADK) that connects directly to this RAG system:

from .tools import query_hr
hr_rag_tool = FunctionTool(func=query_hr)

hr_agent = LlmAgent(
    name="hr_agent",
    model="gemini-3.1-pro-preview",
    description="Specialist in company HR policies and procedures.",
    instruction=(
        "You are a professional HR assistant. Your goal is to answer questions "
        "using ONLY the information retrieved from the 'query_hr' tool. "
        "When calling the 'query_hr' tool, ensure all string arguments are properly formatted as standard JSON strings with double quotes.\n\n"
        "RULES:\n"
        "1. If the tool returns relevant information, summarize it clearly.\n"
        "2. You MUST cite your sources using the format: (Source: [Source Name], Page: [Page Number]).\n"
        "3. If the tool results do not contain the answer, state: 'I'm sorry, I couldn't find that in our HR documents.'\n"
        "4. Do not use outside knowledge or make up facts about company policy."
    ),
    tools=[query_hr],
)

By giving the agent clear instructions and the right tools to search our vector database, it should be able to pull precise answers for users in seconds:

It answers questions reliably, but if I'm being honest, I can't help but feel we're only scratching the surface of the "full" answers that we're looking for.

Next steps

We have a working prototype, but there's still plenty of room to grow. To transform this from a simple RAG system into a high performance engine, our next steps will focus on precision. We'll refine how we process and chunk documents and introduce a reranking layer to our search results to significantly boost the quality of the agent's responses.

Additional learning

If you haven't used Agent Development Kit yet, but would like to learn more, checkout this Codelab: "ADK Crash Course - From Beginner to Expert" (it comes with a link to claim some free GCP credits to get you through the course).