# Building an Imitation Learning Pipeline for Robotic Manipulation

> Source: <https://dev.to/vmodal_ai/building-an-imitation-learning-pipeline-for-robotic-manipulation-4flk>
> Published: 2026-09-07 22:27:24+00:00

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.

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

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.

``` python
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(),
        }
```

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).

``` python
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.

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:

``` python
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.

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

``` python
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}")
```

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:

``` python
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.

| 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 | 

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.

Website: [www.v-modal.com](https://www.v-modal.com)

SDK Flutter: [https://github.com/v-modal/vmodal_sdk_flutter](https://github.com/v-modal/vmodal_sdk_flutter)

SDK Android: [https://github.com/v-modal/vmodal_sdk_android](https://github.com/v-modal/vmodal_sdk_android)

Discord: [https://discord.gg/K72z28KUx](https://discord.gg/K72z28KUx)