Building an Imitation Learning Pipeline for Robotic Manipulation

Building an Imitation Learning Pipeline for Robotic Manipulation

With a validated demonstration dataset in hand, the next step is building the actual training pipeline: turning (observation, action) pairs into a policy that can control the robot on its own. This tutorial covers the end-to-end pipeline — data loading, model architecture choices, training, and evaluation — for behavior-cloning-style imitation learning on manipulation tasks.

Pipeline overview

Raw Episodes -> Preprocessing -> Dataset/Dataloader -> Policy Model -> Training Loop -> Evaluation

Enter fullscreen mode Exit fullscreen mode

Each stage has decisions that materially affect final policy quality, so it’s worth treating this as a real pipeline with clear interfaces, not a single monolithic script.

Step 1: Preprocessing

Before training, normalize your data:

• Action normalization: scale actions to roughly [-1, 1] using dataset statistics (mean/std or min/max). Unnormalized actions with wildly different scales across joints slow down and destabilize training.
• Image preprocessing: resize to a consistent resolution, normalize pixel values, and optionally apply light augmentation (color jitter, random crop) to improve robustness — but avoid augmentations that break task-relevant geometry, like aggressive rotation for tasks sensitive to orientation.
• Action representation: decide between predicting absolute targets vs. delta actions vs. velocity commands. Delta actions relative to the current state are usually easier to learn and more robust to compounding drift at inference time.

import numpy as np

class ActionNormalizer:
    def __init__(self, actions: np.ndarray):
        self.mean = actions.mean(axis=0)
        self.std = actions.std(axis=0) + 1e-6

    def normalize(self, action):
        return (action - self.mean) / self.std

    def denormalize(self, action):
        return action * self.std + self.mean

Enter fullscreen mode Exit fullscreen mode

Step 2: Dataset and dataloader

Structure your dataset so each sample is a short window of context, not just a single frame — most manipulation policies benefit from a few frames of history to disambiguate velocity and resolve partial occlusion.

from torch.utils.data import Dataset

class ManipulationDataset(Dataset):
    def __init__(self, episodes, obs_horizon=2, action_horizon=8):
        self.episodes = episodes
        self.obs_horizon = obs_horizon
        self.action_horizon = action_horizon
        self.index = self._build_index()

    def _build_index(self):
        index = []
        for ep_i, ep in enumerate(self.episodes):
            for t in range(self.obs_horizon, len(ep) - self.action_horizon):
                index.append((ep_i, t))
        return index

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

    def __getitem__(self, idx):
        ep_i, t = self.index[idx]
        ep = self.episodes[ep_i]
        obs_seq = ep[t - self.obs_horizon : t]
        action_seq = ep[t : t + self.action_horizon]
        return {
            "images": [s["observation"]["image"] for s in obs_seq],
            "joint_state": obs_seq[-1]["observation"]["joint_state"],
            "actions": [s["action"] for s in action_seq],
        }

Enter fullscreen mode Exit fullscreen mode

Predicting a short sequence of future actions (action chunking) rather than a single next action, and executing several before re-planning, tends to produce noticeably smoother behavior than pure single-step prediction — this is one of the more impactful architectural choices in recent manipulation policy designs.

Step 3: Model architecture

A reasonable baseline architecture for image-based manipulation:

1. A visual encoder (a small CNN or a pretrained ResNet backbone) processes each camera frame into a feature vector.
2. Visual features are concatenated with proprioceptive state (joint positions/velocities).
3. A sequence model (transformer or a simple MLP with the obs_horizon flattened) maps the combined features to a predicted action chunk.

import torch
import torch.nn as nn

class ManipulationPolicy(nn.Module):
    def __init__(self, visual_encoder, state_dim, action_dim, action_horizon):
        super().__init__()
        self.visual_encoder = visual_encoder
        feature_dim = visual_encoder.output_dim + state_dim
        self.head = nn.Sequential(
            nn.Linear(feature_dim, 512),
            nn.ReLU(),
            nn.Linear(512, action_dim * action_horizon),
        )
        self.action_dim = action_dim
        self.action_horizon = action_horizon

    def forward(self, image, joint_state):
        visual_feat = self.visual_encoder(image)
        combined = torch.cat([visual_feat, joint_state], dim=-1)
        out = self.head(combined)
        return out.view(-1, self.action_horizon, self.action_dim)

Enter fullscreen mode Exit fullscreen mode

For more advanced setups, diffusion-based action heads (predicting the action chunk via denoising) tend to model multimodal human behavior — where the same task can reasonably be solved multiple ways — better than a plain MLP regression head, at the cost of more complex training and slower inference.

Step 4: Training loop

Standard supervised regression training, with mean squared error (or smooth L1) between predicted and demonstrated actions:

def train_epoch(model, dataloader, optimizer, device):
    model.train()
    total_loss = 0.0
    for batch in dataloader:
        images = batch["images"].to(device)
        joint_state = batch["joint_state"].to(device)
        target_actions = batch["actions"].to(device)

        pred_actions = model(images, joint_state)
        loss = nn.functional.smooth_l1_loss(pred_actions, target_actions)

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

    return total_loss / len(dataloader)

Enter fullscreen mode Exit fullscreen mode

A few practical tips:

• Track validation loss on held-out episodes, not held-out frames — random frame-level splits leak information between train and val since adjacent frames are highly correlated.
• Watch for the gap between training loss and real-world performance: low loss doesn’t guarantee good closed-loop behavior, since imitation learning is an open-loop training objective applied to a closed-loop control problem (compounding error is the classic failure mode here).
• Use early stopping based on a small set of real or simulated rollout evaluations if at all possible, not just loss curves.

Step 5: Evaluation

Loss curves only tell part of the story. Evaluate with actual closed-loop rollouts on the robot (or in simulation) and track:

• Task success rate across multiple trials and initial conditions.
• Recovery behavior — does the policy handle small perturbations gracefully, or does any deviation from the training distribution cause it to fail catastrophically?
• Smoothness — jerky, high-frequency action outputs often indicate the policy hasn’t generalized well and is essentially “hunting” between similar training examples.

Where to go from here

This pipeline covers the core behavior-cloning loop. From here, natural extensions include adding more demonstration diversity, experimenting with different action representations, or moving to more advanced input devices like VR controllers to make demonstration collection faster and more natural — which we cover next in this series.

Useful Links

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