{"slug": "building-an-imitation-learning-pipeline-for-robotic-manipulation", "title": "Building an Imitation Learning Pipeline for Robotic Manipulation", "summary": "A developer detailed the construction of an end-to-end imitation learning pipeline for robotic manipulation, covering dataset loading, normalization, model architecture, and training. The tutorial emphasizes action chunking to reduce compounding error and the importance of state normalization, providing code examples for a PyTorch-based policy with a CNN vision encoder.", "body_md": "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.\n\n``` php\n[Raw Episodes] --> [Dataset Loader] --> [Preprocessing] --> [Policy Model]\n                                                                  |\n                                                                  v\n                                                          [Training Loop]\n                                                                  |\n                                                                  v\n                                                    [Checkpoint] --> [Rollout/Eval]\n```\n\nWrap 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.\n\n``` python\nimport torch\nfrom torch.utils.data import Dataset\nimport h5py\nimport glob\n\nclass DemoDataset(Dataset):\n    def __init__(self, data_dir, chunk_size=8):\n        self.files = glob.glob(f\"{data_dir}/*.hdf5\")\n        self.chunk_size = chunk_size\n        self.index = self._build_index()\n\n    def _build_index(self):\n        index = []\n        for file_idx, f in enumerate(self.files):\n            with h5py.File(f, \"r\") as h:\n                length = h[\"actions\"].shape[0]\n                for t in range(length - self.chunk_size):\n                    index.append((file_idx, t))\n        return index\n\n    def __len__(self):\n        return len(self.index)\n\n    def __getitem__(self, idx):\n        file_idx, t = self.index[idx]\n        with h5py.File(self.files[file_idx], \"r\") as h:\n            image = h[\"obs/front_rgb\"][t]\n            state = h[\"obs/joint_positions\"][t]\n            action_chunk = h[\"actions\"][t:t + self.chunk_size]\n\n        return {\n            \"image\": torch.from_numpy(image).permute(2, 0, 1).float() / 255.0,\n            \"state\": torch.from_numpy(state).float(),\n            \"action_chunk\": torch.from_numpy(action_chunk).float(),\n        }\n```\n\nAction 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).\n\n``` python\nimport numpy as np\n\ndef compute_normalization_stats(dataset):\n    all_actions = np.concatenate([dataset[i][\"action_chunk\"].numpy() for i in range(len(dataset))])\n    return {\n        \"action_mean\": all_actions.mean(axis=0),\n        \"action_std\": all_actions.std(axis=0) + 1e-6,\n    }\n\ndef normalize_action(action, stats):\n    return (action - stats[\"action_mean\"]) / stats[\"action_std\"]\n\ndef denormalize_action(action, stats):\n    return action * stats[\"action_std\"] + stats[\"action_mean\"]\n```\n\nSave these stats alongside your checkpoint — you'll need the exact same normalization at inference time.\n\nA 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:\n\n``` python\nimport torch.nn as nn\n\nclass VisionEncoder(nn.Module):\n    def __init__(self, out_dim=256):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Conv2d(3, 32, 5, stride=2), nn.ReLU(),\n            nn.Conv2d(32, 64, 5, stride=2), nn.ReLU(),\n            nn.Conv2d(64, 128, 3, stride=2), nn.ReLU(),\n            nn.AdaptiveAvgPool2d(1),\n        )\n        self.fc = nn.Linear(128, out_dim)\n\n    def forward(self, x):\n        x = self.net(x).flatten(1)\n        return self.fc(x)\n\nclass Policy(nn.Module):\n    def __init__(self, state_dim, action_dim, chunk_size, hidden_dim=256):\n        super().__init__()\n        self.vision = VisionEncoder(out_dim=hidden_dim)\n        self.state_proj = nn.Linear(state_dim, hidden_dim)\n        self.head = nn.Sequential(\n            nn.Linear(hidden_dim * 2, hidden_dim), nn.ReLU(),\n            nn.Linear(hidden_dim, action_dim * chunk_size),\n        )\n        self.chunk_size = chunk_size\n        self.action_dim = action_dim\n\n    def forward(self, image, state):\n        v = self.vision(image)\n        s = self.state_proj(state)\n        combined = torch.cat([v, s], dim=-1)\n        out = self.head(combined)\n        return out.view(-1, self.chunk_size, self.action_dim)\n```\n\nFor 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.\n\nStandard behavior cloning uses MSE loss between predicted and demonstrated actions:\n\n``` python\ndef train(model, dataloader, optimizer, epochs=100, device=\"cuda\"):\n    model.to(device)\n    model.train()\n\n    for epoch in range(epochs):\n        total_loss = 0.0\n        for batch in dataloader:\n            image = batch[\"image\"].to(device)\n            state = batch[\"state\"].to(device)\n            target = batch[\"action_chunk\"].to(device)\n\n            pred = model(image, state)\n            loss = nn.functional.mse_loss(pred, target)\n\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n            total_loss += loss.item()\n\n        print(f\"Epoch {epoch}: loss={total_loss / len(dataloader):.4f}\")\n```\n\nLoss 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:\n\n``` python\ndef rollout(env, model, stats, device=\"cuda\", max_steps=200):\n    obs = env.reset()\n    model.eval()\n\n    for step in range(max_steps):\n        image = preprocess_image(obs[\"front_rgb\"]).to(device)\n        state = torch.from_numpy(obs[\"joint_positions\"]).float().unsqueeze(0).to(device)\n\n        with torch.no_grad():\n            action_chunk = model(image, state)[0].cpu().numpy()\n\n        action = denormalize_action(action_chunk[0], stats)\n        obs, _, done, info = env.step(action)\n\n        if done:\n            break\n\n    return info.get(\"success\", False)\n```\n\nExecuting 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.\n\n| Symptom | Likely Cause | Fix | \n|---|---|---|\n| Training loss low, rollout fails | Distribution shift / compounding error | Add action chunking, DAgger-style correction data, or more diverse demos | \n| Policy freezes near objects | Visual overfitting to background | Augment with random crops/color jitter, add more scene variation | \n| Jerky robot motion | No temporal smoothing | Smooth action chunk execution, reduce control frequency mismatch | \n| Great on training tasks, bad on new positions | Insufficient position diversity in demos | Collect demos with randomized object placement | \n\nThis 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.\n\nWebsite: [www.v-modal.com](https://www.v-modal.com)\n\nSDK Flutter: [https://github.com/v-modal/vmodal_sdk_flutter](https://github.com/v-modal/vmodal_sdk_flutter)\n\nSDK Android: [https://github.com/v-modal/vmodal_sdk_android](https://github.com/v-modal/vmodal_sdk_android)\n\nDiscord: [https://discord.gg/K72z28KUx](https://discord.gg/K72z28KUx)", "url": "https://wpnews.pro/news/building-an-imitation-learning-pipeline-for-robotic-manipulation", "canonical_source": "https://dev.to/vmodal_ai/building-an-imitation-learning-pipeline-for-robotic-manipulation-4flk", "published_at": "2026-09-07 22:27:24+00:00", "updated_at": "2026-09-07 23:00:52.565001+00:00", "lang": "en", "topics": ["robotics", "machine-learning", "developer-tools"], "entities": ["PyTorch", "CNN"], "alternates": {"html": "https://wpnews.pro/news/building-an-imitation-learning-pipeline-for-robotic-manipulation", "markdown": "https://wpnews.pro/news/building-an-imitation-learning-pipeline-for-robotic-manipulation.md", "text": "https://wpnews.pro/news/building-an-imitation-learning-pipeline-for-robotic-manipulation.txt", "jsonld": "https://wpnews.pro/news/building-an-imitation-learning-pipeline-for-robotic-manipulation.jsonld"}}