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 , model architecture, training loop, and evaluation, using the demonstration format from the previous tutorial.
[Raw Episodes] --> [Dataset ] --> [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.
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).
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:
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:
def train(model, data, optimizer, epochs=100, device="cuda"):
model.to(device)
model.train()
for epoch in range(epochs):
total_loss = 0.0
for batch in data:
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(data):.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:
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
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