Imitation Learning for Robots: Teach by Demonstration

Imitation learning replaces hand-written controllers with a policy trained on recorded human demonstrations. You teleoperate the robot, log observations and actions, and fit a neural network that maps one to the other. The method works well for contact-rich manipulation, where reward functions are hard to write and reinforcement learning (RL) needs millions of unsafe trials. This guide covers the math, the data pipeline, the algorithms, and a working ROS2 + PyTorch implementation.

Quick Takeaways

• Behavior cloning (BC) is supervised learning on (observation, action) pairs. It is simple and fast, but it suffers from compounding error.
• DAgger fixes distribution shift by letting an expert relabel states the learner visits.
• ACT and Diffusion Policy predict action chunks. They handle multimodal demonstrations and give smoother trajectories than single-step BC.
• Data quality beats data volume. 50 consistent demonstrations often outperform 500 noisy ones.

Problem Formulation

A demonstration dataset is D = {τ₁, τ₂, …, τₙ}. Each trajectory is τ = (o₁, a₁, o₂, a₂, …, o_T, a_T), where o_t is an observation (camera images, joint positions) and a_t is the expert action.

The goal is a policy π_θ(a | o) that minimizes the divergence from the expert policy π*:

θ* = argmin_θ  E_{(o,a) ~ D} [ L( π_θ(o), a ) ]

L = || π_θ(o) - a ||²          (deterministic, MSE)
L = - log π_θ(a | o)           (stochastic, max likelihood)
L = || π_θ(o) - a ||₁          (L1, preferred in ACT)

The Compounding Error Problem

BC trains on states drawn from the expert’s distribution. At test time the learner drifts into states the expert never visited. Ross et al. showed the error bound grows quadratically with horizon T:

J(π_θ) ≤ J(π*) + T² · ε

Here ε is the per-step imitation error. DAgger reduces the bound to T · ε by training on the learner’s own state distribution.

Hardware and Data Collection Architecture

Demonstration quality sets the ceiling on policy quality. Choose the interface based on the task.

Recommended Sensor Stack

• Wrist RGB camera: 640×480 at 30 FPS for fine alignment.
• Overhead RGB camera: 640×480 at 30 FPS for global scene context.
• Joint encoders: 100–500 Hz joint position and velocity.
• Force/torque sensor (optional): 6-axis at 500 Hz for insertion tasks.

Synchronize all streams with message_filters.ApproximateTimeSynchronizer and a slop of 0.02 s. Log everything with rosbag2.

Recording Demonstrations in ROS2

Record a bag with compressed images to keep disk throughput manageable:

# Record one demonstration episode
ros2 bag record -o demo_001 \
  /joint_states \
  /leader/joint_states \
  /wrist_cam/image_raw/compressed \
  /overhead_cam/image_raw/compressed

The node below records synchronized (observation, action) pairs straight to an HDF5 file. It treats the leader arm’s joint positions as the action and the follower’s state as the observation.

#!/usr/bin/env python3
# demo_recorder.py - Synchronized demo logger for ROS2 Humble/Jazzy
import rclpy, h5py, numpy as np, cv2
from rclpy.node import Node
from sensor_msgs.msg import JointState, CompressedImage
from message_filters import Subscriber, ApproximateTimeSynchronizer

class DemoRecorder(Node):
    def __init__(self):
        super().__init__('demo_recorder')
        # Subscribers: follower state (observation), leader state (action), camera
        self.follower = Subscriber(self, JointState, '/joint_states')
        self.leader   = Subscriber(self, JointState, '/leader/joint_states')
        self.cam      = Subscriber(self, CompressedImage, '/wrist_cam/image_raw/compressed')
        # slop=0.02 s: max timestamp mismatch tolerated between streams
        self.sync = ApproximateTimeSynchronizer(
            [self.follower, self.leader, self.cam], queue_size=10, slop=0.02)
        self.sync.registerCallback(self.cb)
        self.qpos, self.actions, self.images = [], [], []
        self.create_timer(60.0, self.save)  # flush every 60 s; adjust per episode length

    def cb(self, follower, leader, cam):
        self.qpos.append(np.array(follower.position, dtype=np.float32))   # o_t: joint angles (rad)
        self.actions.append(np.array(leader.position, dtype=np.float32))  # a_t: target joint angles (rad)
        img = cv2.imdecode(np.frombuffer(cam.data, np.uint8), cv2.IMREAD_COLOR)
        self.images.append(cv2.resize(img, (224, 224)))                    # downsample for the encoder

    def save(self):
        with h5py.File('episode_000.hdf5', 'w') as f:
            f.create_dataset('qpos', data=np.stack(self.qpos))
            f.create_dataset('action', data=np.stack(self.actions))
            f.create_dataset('images', data=np.stack(self.images), compression='gzip')
        self.get_logger().info(f'Saved {len(self.qpos)} timesteps')

def main():
    rclpy.init(); rclpy.spin(DemoRecorder())

if __name__ == '__main__':
    main()

Data Hygiene Checklist

• Trim idle start and end frames. They teach the policy to stand still.
• Keep demonstration speed consistent. Mixed speeds create multimodal action targets.
• Randomize object start poses over the full workspace, not one corner.
• Record recovery behavior: small mistakes followed by corrections.

Behavior Cloning with Action Chunking

Single-step prediction accumulates error at every control tick. Action chunking predicts k future actions per inference call. This cuts the effective decision horizon by a factor of k.

π_θ(a_{t:t+k} | o_t)     with k = 50 at 50 Hz → 1-second chunks

Temporal ensembling averages overlapping chunk predictions for the current timestep with exponential weights:

a_t = Σ_i  w_i · a_t^(i)  /  Σ_i w_i,     w_i = exp(-m · i)

Here a_t^(i) is the prediction for time t made i steps ago, and m controls how quickly old predictions are discounted (ACT uses m = 0.01).

PyTorch Training Loop

This minimal chunked BC model uses a ResNet18 image encoder and an MLP head. It is a teaching baseline, not a full ACT implementation.

# train_bc.py - Chunked behavior cloning baseline
import torch, torch.nn as nn, torchvision, h5py, numpy as np
from torch.utils.data import Dataset, DataLoader

K = 20          # chunk length: 20 steps at 50 Hz = 0.4 s lookahead
DOF = 7         # 7-DoF arm (e.g., Franka Panda)

class DemoDataset(Dataset):
    def __init__(self, path):
        f = h5py.File(path, 'r')
        self.q, self.a, self.img = f['qpos'][:], f['action'][:], f['images'][:]
    def __len__(self): return len(self.q) - K
    def __getitem__(self, i):
        img = torch.from_numpy(self.img[i]).permute(2, 0, 1).float() / 255.0  # HWC → CHW, [0,1]
        q = torch.from_numpy(self.q[i])
        chunk = torch.from_numpy(self.a[i:i+K])   # shape (K, DOF): future expert actions
        return img, q, chunk

class Policy(nn.Module):
    def __init__(self):
        super().__init__()
        r = torchvision.models.resnet18(weights='IMAGENET1K_V1')
        r.fc = nn.Identity()                      # 512-d image feature
        self.enc = r
        self.head = nn.Sequential(
            nn.Linear(512 + DOF, 512), nn.ReLU(),
            nn.Linear(512, 512), nn.ReLU(),
            nn.Linear(512, K * DOF))              # output flattened action chunk
    def forward(self, img, q):
        z = torch.cat([self.enc(img), q], dim=-1) # fuse vision + proprioception
        return self.head(z).view(-1, K, DOF)

dev = 'cuda' if torch.cuda.is_available() else 'cpu'
model = Policy().to(dev)
opt = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
loader = DataLoader(DemoDataset('episode_000.hdf5'), batch_size=32, shuffle=True)

for epoch in range(100):
    for img, q, chunk in loader:
        img, q, chunk = img.to(dev), q.to(dev), chunk.to(dev)
        loss = nn.functional.l1_loss(model(img, q), chunk)  # L1 loss gives sharper actions than MSE
        opt.zero_grad(); loss.backward(); opt.step()
    print(f'epoch {epoch} loss {loss.item():.4f}')

torch.save(model.state_dict(), 'bc_policy.pt')

Normalize qpos and action tensors with dataset mean and standard deviation before training. Store the statistics beside the checkpoint. Inference must apply the same normalization.

DAgger: Fixing Distribution Shift

DAgger (Dataset Aggregation) iterates between rollout and relabeling:

1. Train π_1 on expert dataset D_0
2. For iteration i = 1..N:
   a. Roll out mixed policy: π_mix = β_i · π* + (1 - β_i) · π_i
   b. Record visited states s_t
   c. Query expert for labels: a*_t = π*(s_t)
   d. Aggregate: D_i = D_{i-1} ∪ {(s_t, a*_t)}
   e. Retrain π_{i+1} on D_i

The mixing coefficient β_i decays across iterations, for example β_i = 0.5^i. Early rounds rely on the expert, and later rounds expose the learner’s own failures.

Practical variant (HG-DAgger): the learner runs autonomously. A human watches and takes over when the robot drifts toward failure. Only the intervention segments are added to the dataset. This keeps the expert safe and avoids labeling states the learner handles correctly.

Modern Architectures: ACT vs. Diffusion Policy

ACT (Action Chunking with Transformers) encodes demonstrations with a conditional VAE. A style variable z captures the demonstrator’s intent, and a transformer decoder outputs the chunk. At test time z = 0, which gives a deterministic mean behavior. ACT trained on 50 demonstrations, 10 minutes of data, achieved high success on bimanual tasks such as threading a zip tie on ALOHA hardware.

Diffusion Policy models the action distribution as a denoising process:

a^{k-1} = α · ( a^k - γ · ε_θ(o, a^k, k) ) + N(0, σ²I)

Here ε_θ predicts the noise added at step k. The network starts from Gaussian noise and iteratively denoises it into an action chunk conditioned on o. It represents multimodal action distributions without mode averaging, which is a failure case for MSE-trained BC. When an object can be grasped from the left or the right, MSE regression averages both into a motion that reaches neither side. Diffusion samples one mode.

For production work, use the LeRobot library from Hugging Face. It ships ACT, Diffusion Policy, and dataset tooling:

pip install lerobot
python lerobot/scripts/train.py \
  policy=act \
  dataset_repo_id=your_user/your_dataset \
  training.batch_size=8 \
  training.num_workers=4

Deploying the Policy as a ROS2 Node

The inference node runs the policy at a low rate and publishes joint targets at a high rate with temporal ensembling:

#!/usr/bin/env python3
# policy_node.py - Chunked inference with temporal ensembling
import rclpy, torch, numpy as np
from rclpy.node import Node
from sensor_msgs.msg import JointState
from std_msgs.msg import Float64MultiArray
from train_bc import Policy, K, DOF

class PolicyNode(Node):
    def __init__(self):
        super().__init__('policy_node')
        self.model = Policy(); self.model.load_state_dict(torch.load('bc_policy.pt')); self.model.eval()
        self.pub = self.create_publisher(Float64MultiArray, '/joint_position_controller/commands', 10)
        self.create_subscription(JointState, '/joint_states', self.on_state, 10)
        self.buffer = {}                      # timestep → list of predicted actions
        self.t = 0
        self.create_timer(0.02, self.step)    # 50 Hz control loop
        self.q = np.zeros(DOF, dtype=np.float32)
        self.img = torch.zeros(1, 3, 224, 224)  # replace with a live camera subscriber

    def on_state(self, msg): self.q = np.array(msg.position[:DOF], dtype=np.float32)

    def step(self):
        with torch.no_grad():
            chunk = self.model(self.img, torch.from_numpy(self.q)[None])[0].numpy()  # (K, DOF)
        for i in range(K):                    # register this chunk for the next K timesteps
            self.buffer.setdefault(self.t + i, []).append(chunk[i])
        preds = self.buffer.pop(self.t)       # all predictions targeting the current step
        w = np.exp(-0.01 * np.arange(len(preds)))  # m = 0.01 exponential weights
        w = w[::-1] / w.sum()                 # newest prediction gets the largest weight
        action = np.sum(np.array(preds) * w[:, None], axis=0)
        self.pub.publish(Float64MultiArray(data=action.tolist()))
        self.t += 1

def main():
    rclpy.init(); rclpy.spin(PolicyNode())

if __name__ == '__main__':
    main()

Run the ensemble only on smooth tasks. For discrete events such as grasp closure, ACT reports better results with chunk execution and no ensembling. Test both on your task.

Real-World Workflow: Training a Pick-and-Place Policy

Follow this pipeline for a 6-DoF arm placing a cube into a bin:

1. Define success. The cube rests inside the bin within 10 mm of the center.
2. Collect 60 demonstrations with leader-follower teleoperation. Randomize the cube pose over a 30 cm × 30 cm region.
3. Filter. Drop episodes longer than 1.5× the median length. Those usually contain hesitation or regrasps.
4. Train ACT with chunk size 50 and 2,000 epochs. Hold out 10% of episodes for validation loss.
5. Evaluate 20 rollouts on novel cube poses. Record success rate and failure modes.
6. Run HG-DAgger. Collect 20 intervention segments on observed failures, merge them into the dataset, and retrain.
7. Freeze and benchmark. Export to ONNX or TensorRT and verify end-to-end latency under 30 ms.

Typical outcome: the first-pass success rate sits around 60–75%. One or two DAgger rounds push it above 90% on in-distribution poses.

Failure Modes and Fixes

Testing and Safety

• Validate in simulation first with MuJoCo or Isaac Sim, using scripted demonstrations.
• Clamp joint targets to hardware limits in the controller, never only in the policy.
• Limit the per-tick delta between consecutive commands, for example 0.05 rad.
• Wire a hardware e-stop that is independent of the inference stack.
• Log every deployed rollout. Failed runs are the most valuable data for the next DAgger round.

FAQ

What is the difference between imitation learning and reinforcement learning?

Imitation learning fits a policy to expert demonstrations with supervised loss. It needs no reward function. Reinforcement learning optimizes a reward through trial and error and usually needs far more interaction. Imitation learning is faster to deploy on real hardware, but it cannot exceed the demonstrator without extra fine-tuning.

How many demonstrations does a robot need to learn a task?

Simple tasks need 50–100 demonstrations with modern methods such as ACT or Diffusion Policy. Tasks with wide object variation can need 200–1,000. Consistency matters more than count, so keep a fixed strategy and speed across demonstrations.

What is behavior cloning and why does it fail?

Behavior cloning is supervised regression from observations to expert actions. It fails because small errors push the robot into states absent from the training data, where predictions degrade further. The resulting error grows quadratically with task horizon. DAgger, action chunking, and larger state coverage reduce the effect.

Which is better for robot manipulation: ACT or Diffusion Policy?

ACT trains faster and runs with lower latency, which suits bimanual tasks and edge hardware. Diffusion Policy captures multimodal action distributions more reliably and performs well on contact-rich tasks, but inference is slower. Start with ACT and move to Diffusion Policy if mode averaging appears.