Building an Imitation Learning Pipeline for Robotic Manipulation

You've collected demonstrations — now it's time to turn them into a working policy. This tutorial walks through building an end-to-end imitation learning pipeline: dataset loading, model architecture, training loop, and evaluation, using the demonstration format from the previous tutorial.

Pipeline Overview

[Raw Episodes] --> [Dataset Loader] --> [Preprocessing] --> [Policy Model]
                                                                  |
                                                                  v
                                                          [Training Loop]
                                                                  |
                                                                  v
                                                    [Checkpoint] --> [Rollout/Eval]

Step 1: Dataset Loader

Wrap your recorded episodes in a PyTorch Dataset that returns (observation, action) pairs, optionally as short action chunks rather than single steps — chunking (predicting several future actions at once) is standard in modern imitation learning and reduces compounding error.

import torch
from torch.utils.data import Dataset
import h5py
import glob

class DemoDataset(Dataset):
    def __init__(self, data_dir, chunk_size=8):
        self.files = glob.glob(f"{data_dir}/*.hdf5")
        self.chunk_size = chunk_size
        self.index = self._build_index()

    def _build_index(self):
        index = []
        for file_idx, f in enumerate(self.files):
            with h5py.File(f, "r") as h:
                length = h["actions"].shape[0]
                for t in range(length - self.chunk_size):
                    index.append((file_idx, t))
        return index

    def __len__(self):
        return len(self.index)

    def __getitem__(self, idx):
        file_idx, t = self.index[idx]
        with h5py.File(self.files[file_idx], "r") as h:
            image = h["obs/front_rgb"][t]
            state = h["obs/joint_positions"][t]
            action_chunk = h["actions"][t:t + self.chunk_size]

        return {
            "image": torch.from_numpy(image).permute(2, 0, 1).float() / 255.0,
            "state": torch.from_numpy(state).float(),
            "action_chunk": torch.from_numpy(action_chunk).float(),
        }

Step 2: Preprocessing and Normalization

Action and state normalization is one of the most impactful — and most skipped — steps. Compute statistics over your full dataset and normalize to zero mean / unit variance (or a fixed range).

import numpy as np

def compute_normalization_stats(dataset):
    all_actions = np.concatenate([dataset[i]["action_chunk"].numpy() for i in range(len(dataset))])
    return {
        "action_mean": all_actions.mean(axis=0),
        "action_std": all_actions.std(axis=0) + 1e-6,
    }

def normalize_action(action, stats):
    return (action - stats["action_mean"]) / stats["action_std"]

def denormalize_action(action, stats):
    return action * stats["action_std"] + stats["action_mean"]

Save these stats alongside your checkpoint — you'll need the exact same normalization at inference time.

Step 3: Policy Architecture

A common and effective baseline is a CNN vision encoder feeding into an MLP or transformer head that predicts an action chunk. Here's a simple version:

import torch.nn as nn

class VisionEncoder(nn.Module):
    def __init__(self, out_dim=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 32, 5, stride=2), nn.ReLU(),
            nn.Conv2d(32, 64, 5, stride=2), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.fc = nn.Linear(128, out_dim)

    def forward(self, x):
        x = self.net(x).flatten(1)
        return self.fc(x)

class Policy(nn.Module):
    def __init__(self, state_dim, action_dim, chunk_size, hidden_dim=256):
        super().__init__()
        self.vision = VisionEncoder(out_dim=hidden_dim)
        self.state_proj = nn.Linear(state_dim, hidden_dim)
        self.head = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, action_dim * chunk_size),
        )
        self.chunk_size = chunk_size
        self.action_dim = action_dim

    def forward(self, image, state):
        v = self.vision(image)
        s = self.state_proj(state)
        combined = torch.cat([v, s], dim=-1)
        out = self.head(combined)
        return out.view(-1, self.chunk_size, self.action_dim)

For more capable policies, consider swapping the head for a small transformer decoder or a diffusion head (as in Diffusion Policy) — but this MLP baseline is a good place to validate your whole pipeline before adding complexity.

Step 4: Training Loop

Standard behavior cloning uses MSE loss between predicted and demonstrated actions:

def train(model, dataloader, optimizer, epochs=100, device="cuda"):
    model.to(device)
    model.train()

    for epoch in range(epochs):
        total_loss = 0.0
        for batch in dataloader:
            image = batch["image"].to(device)
            state = batch["state"].to(device)
            target = batch["action_chunk"].to(device)

            pred = model(image, state)
            loss = nn.functional.mse_loss(pred, target)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        print(f"Epoch {epoch}: loss={total_loss / len(dataloader):.4f}")

Step 5: Evaluation

Loss curves only tell part of the story — the real test is closed-loop rollout on the robot (or in simulation). Structure your eval loop to reuse the exact same normalization and observation pipeline as training:

def rollout(env, model, stats, device="cuda", max_steps=200):
    obs = env.reset()
    model.eval()

    for step in range(max_steps):
        image = preprocess_image(obs["front_rgb"]).to(device)
        state = torch.from_numpy(obs["joint_positions"]).float().unsqueeze(0).to(device)

        with torch.no_grad():
            action_chunk = model(image, state)[0].cpu().numpy()

        action = denormalize_action(action_chunk[0], stats)
        obs, _, done, info = env.step(action)

        if done:
            break

    return info.get("success", False)

Executing only the first action of a predicted chunk, then re-planning, is a common and robust strategy (receding-horizon control) — it limits drift from small model errors.

Common Failure Modes and Fixes

Symptom	Likely Cause	Fix
Training loss low, rollout fails	Distribution shift / compounding error	Add action chunking, DAgger-style correction data, or more diverse demos
Policy freezes near objects	Visual overfitting to background	Augment with random crops/color jitter, add more scene variation
Jerky robot motion	No temporal smoothing	Smooth action chunk execution, reduce control frequency mismatch
Great on training tasks, bad on new positions	Insufficient position diversity in demos	Collect demos with randomized object placement

What's Next

This baseline pipeline gets you to a working policy quickly. From here, natural extensions include diffusion-based action heads, transformer backbones (e.g., ACT-style architectures), and multi-task training across several manipulation tasks.

Useful Links

Website: www.v-modal.com
SDK Flutter: https://github.com/v-modal/vmodal_sdk_flutter
SDK Android: https://github.com/v-modal/vmodal_sdk_android
Discord: https://discord.gg/K72z28KUx