# Imitation Learning for Robotics: A Journey from Scratch to Advanced Flow Matching Policy > Source: > Published: 2026-09-11 22:01:02+00:00 Imitation Learning (IL) is the most promising training paradigm in 2026 for robotic manipulation. However, there are limited resources for building policies from scratch. We created this guide to help you get started with vision-based imitation learning policies. We aim to build intuition by starting small and gradually explaining the core concepts in the field, on the journey to advanced policies. All code is in PyTorch and runs on a single GPU. The full repository — including model implementations, training scripts, and pre-trained checkpoints — is [available here](https://github.com/nizan-mashall/imitation-learning-tutorial.git). Choosing the right training paradigm is highly dependent on the task and on the goal of the robot. **Vision-Language Action (VLA) models**, which are a common and trending term in the robotics community, are actually an imitation learning methods, conditioned by language and trained on a vast amount of data, contributing to their emerging capability to generalize better for language conditions. For example, by conditioning on language, a VLA trained only on a red cube can generalize to picking a green cube — simply by changing the instruction from *“pick the red cube”* to *“pick the green cube”*. While training VLAs from scratch requires enormous compute (OpenVLA [5] — 64 A100 GPUs for 14 days), designing and training your own imitation learning policy is far more affordable — and a great place to build hands-on intuition for the same underlying principles. In robotics, data acquisition is quite hard. Collecting real robot records requires hardware and is highly sensitive to sensor noise and differences in robot embodiment. To avoid the need for hardware and focus only on the learning, we will use **MuJoCo** simulation. MuJoCo is one of the most popular physics simulators for robotics research. Unlike kinematic simulation, which only demonstrates how a robot moves geometrically, physics simulation computes the actual physical forces acting on the robot — such as gravity, friction, and contact — producing estimates of sensor readings. On top of MuJoCo, we will use the **RoboSuite** package. RoboSuite provides multiple task environments for robot learning, including sensor outputs, success indicators, and a uniform benchmark for model evaluation. We will test our model on the **Lift** task, where a robotic arm needs to pick up a red cube. The cube and the arm’s initial position are randomized at every episode, making it harder for the model to memorize a fixed trajectory — forcing it to actually learn. Imitation learning is based on demonstrations — raw trajectories showing the correct way to accomplish a task. Demonstrations can be collected manually or generated by a trained policy. While both methods are valid, researchers have shown that learning from human demonstrations is significantly harder. However, for many applications, learning from a pre-trained policy somewhat defeats the purpose. For this reason, we use **200 prerecorded human demonstrations** provided by the Robomimic [6] package. ***Key insight:*** *One of the major challenges in IL is that there is no single correct trajectory to accomplish a task. Two demonstrations that both succeed may look completely different, which can confuse the model during training — a challenge we will address explicitly when we introduce Flow Matching.* ***Key insight:*** *While one could consider including “bad” demonstrations — showing how a task should not be performed — to expand the data distribution, this has been shown to significantly reduce policy performance. A widely used method to expand the data distribution is to intervene during execution when the policy makes a mistake, and add those corrected trajectories to the training dataset, as proposed in DAgger [7] and demonstrated in BC-Z [8] for imitation learning.* Code for downloading Robomimic demonstrations and rendering robot observations: ``` python import osimport jsonimport subprocessimport h5pyimport numpy as npimport robosuite as suiteREPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))DATASETS_DIR = os.path.join(REPO_ROOT, "datasets")input_path = os.path.join(DATASETS_DIR, "lift", "ph", "low_dim_v141.hdf5")output_path = os.path.join(DATASETS_DIR, "lift_dataset.hdf5")os.makedirs(DATASETS_DIR, exist_ok=True)if not os.path.exists(input_path): print("Downloading dataset...") subprocess.run([ "python", "-m", "robomimic.scripts.download_datasets", "--tasks", "lift", "--dataset_types", "ph", "--download_dir", DATASETS_DIR ], check=True) print("Download complete.")f = h5py.File(input_path, 'r')env_args = json.loads(f['data'].attrs['env_args'])env = suite.make( env_args['env_name'], robots=env_args['env_kwargs']['robots'], has_renderer=False, has_offscreen_renderer=True, use_camera_obs=True, camera_names=["agentview", "robot0_eye_in_hand"], camera_heights=256, camera_widths=256,)out_f = h5py.File(output_path, 'w')out_data = out_f.create_group('data')demos = list(f['data'].keys())#[:200]print(f"Processing {len(demos)} demos...")for i, demo_key in enumerate(demos): demo = f['data'][demo_key] states = demo['states'][:] actions = demo['actions'][:] env.reset() env.sim.set_state_from_flattened(states[0]) env.sim.forward() agentview_images = [] wrist_images = [] eef_pos = [] eef_quat = [] gripper_qpos = [] for t in range(len(actions)): env.sim.set_state_from_flattened(states[t]) env.sim.forward() obs = env._get_observations(force_update=True) agentview_images.append(obs['agentview_image'].copy()) wrist_images.append(obs['robot0_eye_in_hand_image'].copy()) eef_pos.append(obs['robot0_eef_pos'].copy()) eef_quat.append(obs['robot0_eef_quat'].copy()) gripper_qpos.append(obs['robot0_gripper_qpos'].copy()) out_demo = out_data.create_group(demo_key) out_demo.create_dataset('actions', data=actions) obs_group = out_demo.create_group('obs') obs_group.create_dataset('agentview_image', data=np.array(agentview_images)) obs_group.create_dataset('robot0_eye_in_hand_image', data=np.array(wrist_images)) obs_group.create_dataset('robot0_eef_pos', data=np.array(eef_pos)) obs_group.create_dataset('robot0_eef_quat', data=np.array(eef_quat)) obs_group.create_dataset('robot0_gripper_qpos', data=np.array(gripper_qpos)) if (i+1) % 10 == 0: print(f"Processed {i+1}/{len(demos)} demos")f.close()out_f.close()env.close()print("Done! Saved to", output_path) ``` After this step, each trajectory contains: Now we have 200 trajectories, with about 100 steps each, including images from 2 cameras, a vector of the current state, and a vector of the next action. Naively, we could train the policy without images. It will converge fast and could perform complex trajectories very well in open-loop settings. Unfortunately, using only the proprioception (self-state) vector results in a non-reactive — blind policy. The way classic explicit policies use images is by applying an object detector or 3D position estimator to set an explicit dynamic goal. In imitation learning, we aim to learn from data that is as raw and rich as possible. So instead, we use pretrained DinoV2[1], one of the best generalist feature extractors. ***Important note:*** *Unlike traditional feature extractors such as ORB, or CNN-based models that learn kernels to detect edges and textures, DINOv2 is based on a Vision Transformer (ViT) trained in an unsupervised manner to capture high-level semantic concepts — making it significantly better for generalization across scenes and objects.* DINOv2 takes an image, splits it into patches (typically 14×14 pixels), flattens them, and treats each one as a token for the Vision Transformer. After processing, it outputs: We do not train DINOv2 and use it purely as a frozen feature extractor. ***Important note:*** *While using a frozen DINOv2 saves enormous compute by leveraging representations learned from millions of images, several papers have shown that fine-tuning it after a warmup phase can further improve performance.* After processing the images with DINOv2, we can build the state vector by concatenating: **Total: 777-dim observation vector** — a rich summary of what the robot sees and where it is. Every policy’s goal is to predict the next action the robot should take. A common simplification is to learn a **Markovian policy**, which assumes that the next action is a function of only the current state and observation — the history doesn’t matter. ***Key insight:*** *While this is a significant simplification, most VLAs and imitation learning methods are based on it and achieve very impressive results. However, it makes the policy sensitive in some cases — for example, a deterministic Markovian policy cannot recover from a complete stop, and a stochastic Markovian policy might recover but cannot learn to wait for synchronization with other agents or events.* A natural way to map our combined observation and state vector to a 7-dimensional action vector is with a **Multi-Layer Perceptron (MLP)**. We will use DINOv2 as a frozen backbone and a simple MLP head for regression. ``` python import torchimport torch.nn as nnimport torch.nn.functional as Fimport numpy as npclass MlpPolicy(nn.Module): def __init__(self, output_dim, hidden_dims = [512, 512, 256]): super().__init__() self.dino = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14') for param in self.dino.parameters(): param.requires_grad = False dino_CLS_dim = 384 nums_of_cameras = 2 robot_DoF = 9 input_dim = dino_CLS_dim * nums_of_cameras + robot_DoF layers = [] prev_dim = input_dim for hidden_dim in hidden_dims: layers.append(nn.Dropout(0.1)) layers.append(nn.Linear(prev_dim, hidden_dim)) layers.append(nn.LayerNorm(hidden_dim)) layers.append(nn.ReLU()) prev_dim = hidden_dim layers.append(nn.Linear(prev_dim, output_dim)) self.network = nn.Sequential(*layers) def _extract(self, image): out = self.dino.forward_features(image) cls = out['x_norm_clstoken'] return cls def forward(self, x, image1, image2): with torch.no_grad(): image1_features = self._extract(image1) image2_features = self._extract(image2) input_features = torch.cat([x, image1_features, image2_features], dim=1) action = self.network(input_features) return action ``` The training paradigm is similar to supervised learning for regression. The loss is calculated using **L1 loss** — a linear connection between the prediction error and the ground truth action. Although the end-effector is binary (open/close only), treating it as continuous simplifies the architecture with one less hyperparameter to tune. Alternatively, it could be treated as a binary cross-entropy classification problem. ``` ### Main Training Loop model = DinoMlp(output_dim=7, hidden_dims=hidden_dims).to(device) optimizer = torch.optim.AdamW( [p for p in model.parameters() if p.requires_grad], lr=lr) criterion = nn.L1Loss() for epoch in range(1, epochs + 1): model.train() train_loss = 0.0 for images1, images2, states, actions in train_loader: images1 = images1.to(device) images2 = images2.to(device) states = states.to(device) actions = actions.to(device) preds = model(states, images1, images2) loss = criterion(preds, actions) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() train_loss += loss.item() ``` Evaluation can be split in two: evaluating the training process and evaluating the model’s performance. For training evaluation, the most important values to monitor are the training loss and validation loss on trajectories held out from training. Additionally, tracking loss per action dimension is very informative — it can reveal whether training has plateaued overall, or whether some dimensions are still improving while others have converged. For model evaluation, the most important metric is the **success rate**. RoboSuite provides a dedicated success flag indicating whether the cube was lifted. By running multiple rollouts, we can estimate the success rate reliably. An additional interesting metric is **sensitivity to noise** — by increasing the initial position randomization of the robotic arm, we can measure how robustness degrades with perturbation. Furthermore, training with different dataset sizes and measuring the effect on success rate gives valuable insight into the data efficiency of the policy. The MLP policy performs well on simple tasks such as Lift, but is limited in scaling capabilities. For example, using DINOv2 per-patch encodings (instead of the CLS token) is known to dramatically improve model capabilities due to better spatial representation — however, concatenating 256 patches of 384-dim each results in an unreasonably large input layer. Additionally, MLP-based models have no memory mechanism, unlike RNNs or LSTMs, which is an important limitation for sequential decision making. Therefore, we shift from a simple MLP to a **Transformer-based policy**. Without diving too deep, Transformers are sequential models like RNNs and LSTMs. They take a series of vectors of the same dimension and learn — in the **encoder** — the connections between those vectors. Those vectors can represent text, images, state, or any other modality. Then, in the **decoder**, the model learns to output an answer to a specific query based on the encoded representations. Through backpropagation, the model learns the connections between the representations, the queries, and the output — in our case, supervised by L1 loss. ***Important note:*** *Transformers have no inherent sense of order — they do not know which position each token occupies. To address this, we use learned positional embeddings. Without them, the transformer cannot learn the relationship between different cameras and the robot state.* ``` python import torchimport torch.nn as nn class TransformerPolicy(nn.Module): def __init__( self, action_dim=7, state_dim=9, hidden_dim=256, n_heads=8, n_enc_layers=2, n_dec_layers=2, horizon=6, ): super().__init__() self.action_dim = action_dim self.dino = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14') for p in self.dino.parameters(): p.requires_grad = False dino_dim = 384 self.action_horizon = horizon self.patch_proj = nn.Linear(dino_dim, hidden_dim) self.state_proj = nn.Linear(state_dim, hidden_dim) self.num_patches = 256 self.pos_embed_wrist0 = nn.Parameter(torch.randn(1, self.num_patches, hidden_dim) * 0.02) self.pos_embed_agent = nn.Parameter(torch.randn(1, self.num_patches, hidden_dim) * 0.02) encoder_layer = nn.TransformerEncoderLayer( d_model=hidden_dim, nhead=n_heads, dim_feedforward=hidden_dim * 4, dropout=0.1, batch_first=True, ) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_enc_layers) decoder_layer = nn.TransformerDecoderLayer( d_model=hidden_dim, nhead=n_heads, dim_feedforward=hidden_dim * 4, dropout=0.1, batch_first=True, ) self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=n_dec_layers) self.action_query = nn.Parameter(torch.randn(1, self.action_horizon, hidden_dim) * 0.02) self.action_head = nn.Linear(hidden_dim, action_dim) def _extract_tokens(self, image): with torch.no_grad(): out = self.dino.forward_features(image) patches = out['x_norm_patchtokens'] return self.patch_proj(patches) def forward(self, state, image1, image2): B = image2.shape[0] tok1 = self._extract_tokens(image1) tok2 = self._extract_tokens(image2) tok1 = tok1 + self.pos_embed_wrist0 tok2 = tok2 + self.pos_embed_agent state_tok = self.state_proj(state).unsqueeze(1) memory = torch.cat([tok1, tok2, state_tok], dim=1) memory = self.encoder(memory) B = memory.shape[0] action_queries = self.action_query.expand(B, -1, -1) decoded = self.decoder(tgt=action_queries, memory=memory) actions = self.action_head(decoded) return actions ``` Our model is now capable of performing the Lift task. However, when deploying on a real robot, an additional property becomes crucial — **inference frequency**. To achieve smooth trajectories, robots should receive commands at approximately **20 Hz** — one command every 50 milliseconds. Otherwise, the robot will move jerkily, which directly hurts performance. To address this, we modify the model to predict not just the next action, but the next **K actions** — a technique known as **action chunking**. This decouples the low-level command frequency from the inference time. Predicting K steps ahead also addresses an additional challenge: temporal commitment. If the policy predicts only the immediate next step, it cannot account for dynamic properties such as momentum or future direction. By predicting K steps into the future, the model no longer just asks *“Where do I move next?”* but rather *“How does my current movement serve the trajectory over the next K steps?”* To utilize future predictions at inference time, we use an **exponential moving average (EMA)** to blend overlapping predictions — giving higher weight to more recent ones. ``` for step in range(NUM_STEPS): ... with torch.no_grad(): action_chunk = model(states, images1, images2) chunk_np = action_chunk.squeeze(0).cpu().numpy() chunk_buffer.appendleft((chunk_np, step)) proposed_actions = [] weights = [] for age, (chunk, origin_step) in enumerate(chunk_buffer): offset = step - origin_step proposed_actions.append(chunk[offset]) weights.append(np.exp(-0.1 * age)) weights = np.array(weights) weights = weights / weights.sum() action_np = np.average(proposed_actions, axis=0, weights=weights) ``` When deterministic policies, such as the ones we explored so far, encounter two demonstrations in the training dataset where the same observation O and state S lead to different actions, the model cannot separate those cases and instead learns to average them. For example, if there is an obstacle with two valid ways to go around it — one from the left and one from the right — the deterministic model will average the options and might try to go through the middle, crashing into the obstacle. Instead, some policies collapse to a single modality and always use it — for example, always going from the left side. Flow Matching evolved from the family of diffusion models and is based on randomization to avoid “collapsing to a single modality” (always picking the left). It turns our deterministic policy into a stochastic one, which handles different valid actions for the same S and O by learning a distribution instead of a direct function. The core idea of Flow Matching is to learn a velocity map from noise to the target action, conditioned on the observation and time. The noise x₀ is sampled from a Gaussian distribution, and the time t is sampled from a uniform distribution. xₜ is a linear interpolation between the noise and the goal action. The policy learns to predict **v_θ** — the velocity field — based on the current interpolated action xₜ, the sampled time t, and the observation, trained by MSE loss. ``` for images1, images2, states, actions in train_loader: x1 = actions B = x1.shape[0] x0 = torch.randn_like(x1) t = torch.rand(B, device=device) t_expanded = t.view(B, 1, 1) x_t = x0 * (1-t_expanded) + x1 * t_expanded target_velocity = x1 - x0 velocity_pred = model(x_t, t, states, images1, images2) loss = criterion(velocity_pred, target_velocity).mean() optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() ``` After training, action prediction becomes a denoising process — inspired by the kinematic law for constant velocity: In practice, this is not done in a single denoising step, but rather as a loop with incremental steps in the velocity direction. For denoising, x₀ is initialized as during training, but time starts from t = 0 and gradually increases to t = nΔt. Flow Matching requires significantly fewer denoising steps than diffusion models — typically around 10 steps for Flow Matching versus ~1000 for diffusion models — making inference much faster. ``` python @torch.no_grad() def sample(self, state, image1, image2, steps = 10): B = image1.shape[0] x_t = torch.randn(B, self.action_horizon, self.action_dim, device=device) dt = 1.0 / steps for step in range(steps): t_val = step / steps t = torch.full((B,), t_val, device=device, dtype=torch.float32) v_pred = self.forward(x_t, t, state, image1, image2) x_t = x_t + v_pred * dt actions = x_t return actions ``` After training the policy, we can visualize the denoising process — how the learned velocity field guides random noise samples (×) toward a tight cluster of predictions (●) that closely match the ground truth action (★) to be sent as commands to the robot. By the end of this tutorial, we have built three increasingly capable imitation learning policies — from a simple MLP to a Transformer with action chunking, all the way to a Flow Matching policy — and evaluated them on a real robotic manipulation benchmark. The simplest next step I recommend is to try the policies on other RoboSuite tasks that have Robomimic records, and work to identify and understand each policy’s limitations in practice. For those interested in going deeper, a genuinely important open problem is **non-Markovian, memory-based policies**. Every policy we built here acts on the current observation alone — no memory of what came before. **World Action Models (WAMs) [9]** are one of the most promising directions to address this, and are currently a trending research topic in robotics. For those who want to continue the journey toward full VLAs, I recommend fine-tuning an existing VLA rather than building from scratch — the GPU requirements for training from scratch are prohibitive. **SmolVLA [10]** is a great place to start. For those who want to move from simulation to a real robot, [**LeRobot**](https://github.com/huggingface/lerobot) by HuggingFace is an excellent next step. It provides a well-documented robot arm kit, costing around $200 (SO-101), with ready-made IL policies. [1] Oquab, Maxime, et al. “Dinov2: Learning robust visual features without supervision.” *arXiv preprint arXiv:2304.07193* (2023). [2] Zhao, Tony Z., et al. “Learning fine-grained bimanual manipulation with low-cost hardware.” *arXiv preprint arXiv:2304.13705* (2023). [3] Chi, Cheng, et al. “Diffusion policy: Visuomotor policy learning via action diffusion.” *The International Journal of Robotics Research* 44.10–11 (2025): 1684–1704. [4] Carion, Nicolas, et al. “End-to-end object detection with transformers.” *European conference on computer vision*. Cham: Springer International Publishing, 2020. [5] Kim, Moo Jin, et al. “Openvla: An open-source vision-language-action model.” *arXiv preprint arXiv:2406.09246* (2024). [6] Mandlekar, Ajay, et al. “What matters in learning from offline human demonstrations for robot manipulation.” *arXiv preprint arXiv:2108.03298* (2021). [7] Ross, Stéphane, Geoffrey Gordon, and Drew Bagnell. “A reduction of imitation learning and structured prediction to no-regret online learning.” Proceedings of the fourteenth international conference on artificial intelligence and statistics. JMLR Workshop and Conference Proceedings, 2011. [8] Jang, Eric, et al. “Bc-z: Zero-shot task generalization with robotic imitation learning.” conference on Robot Learning. PMLR, 2022. [9] Ye, Seonghyeon, et al. “World action models are zero-shot policies.” *arXiv preprint arXiv:2602.15922* (2026). [10] Shukor, Mustafa, et al. “Smolvla: A vision-language-action model for affordable and efficient robotics.” *arXiv preprint arXiv:2506.01844* (2025). [Imitation Learning for Robotics: A Journey from Scratch to Advanced Flow Matching Policy](https://pub.towardsai.net/imitation-learning-for-robotics-a-journey-from-scratch-to-advanced-flow-matching-policy-a6b2f97883d2) was originally published in [Towards AI](https://pub.towardsai.net) on Medium, where people are continuing the conversation by highlighting and responding to this story.