Skip to main content

Command Palette

Search for a command to run...

Understanding UMI by Running It: How Robot Imitation Learning Works

A Practical Walkthrough of UMI Data Collection, Processing, and Training

Updated
27 min readView as Markdown
Understanding UMI by Running It: How Robot Imitation Learning Works
Y
Factory Automation | Manufacturing Engineering | Validation for Medical Device

This article explains UMI: Universal Manipulation Interface, a framework for collecting robot demonstration data and training manipulation policies. We will follow the official pipeline from raw human demonstrations to model training:

A human demonstrates a task with a handheld gripper
          ↓
Video and IMU data are recorded in a GoPro MP4 file
          ↓
SLAM reconstructs the 6-DoF end-effector trajectory
Markers are used to recover the gripper width
          ↓
Images and actions are time-synchronized and stored in Zarr
          ↓
Short time-series windows are extracted
          ↓
A Diffusion Policy learns sequences of future actions

I also loaded the official cup_in_the_wild.zarr.zip dataset on a PC with a single RTX 4090 and ran 1,000 training steps with a reduced configuration. This is not a full reproduction of the paper's compute setup or real-robot success rate. Instead, the goal is to see exactly how UMI data reaches the learning system by inspecting the actual images, trajectories, gripper widths, and training loss.


1. What Is UMI?

Universal Manipulation Interface, or UMI, is a handheld gripper-based data collection and policy learning framework that turns human demonstrations into data a robot can learn from. It was proposed by a research team from Stanford University, Columbia University, and Toyota Research Institute.

UMI project page

Two common ways to teach a robot a task are:

Method What you do Main challenge
Robot teleoperation A human controls a robot to collect demonstration data You need the robot on site. Setup is heavy for each environment, and controlling the robot can be difficult.
Learning from human videos Use videos of people performing the task Human hands and robot grippers look and move very differently. Ordinary videos also lack the action data a robot needs.

UMI aims for a middle ground. Instead of demonstrating with bare hands, a person uses a handheld device shaped more like a robot gripper. This reduces the visual and kinematic gap between the human demonstration and the robot, while still allowing data to be collected without bringing a robot into every environment.

An important point about the UMI paper is that it is not simply about building a convenient handheld gripper. UMI jointly designs three parts of the system:

  • Data collection device: a handheld gripper with a GoPro

  • Data conversion pipeline: reconstruction of SLAM trajectories, gripper width, and training data from MP4 files

  • Policy interface: a Diffusion Policy that uses RGB images, relative end-effector poses, and gripper width

In other words, UMI combines hardware, data processing, and the learning method into one integrated system.


2. A Robot Needs Both “Observations” and “Actions” to Learn

Training data for robot imitation learning is fundamentally a set of paired observations and actions recorded at matching timestamps.

Observations:
- What the camera can see
- The current end-effector pose
- How far the gripper is open

Actions:
- Which direction the end effector should move next
- How the end effector should rotate
- How much the gripper should open or close

Suppose we record a video of someone placing a cup on a saucer. The video shows that the cup was lifted, but it does not directly tell us the 6-DoF trajectory a robot would need to reproduce the motion. On the other hand, end-effector coordinates alone do not tell us which cup was being moved to which saucer, or what the scene looked like at the time.

What we need is this correspondence:

The image visible at that moment
        ↕ time synchronization
The end-effector position, orientation, and gripper width at that moment

UMI is designed to collect this pairing using only the handheld gripper. On the hardware side, it captures images that resemble what the robot will see during execution. In post-processing, it reconstructs actions in a form the robot can use.


3. UMI Hardware: A Handheld Gripper with a GoPro

UMI's demonstration device is a 3D-printed handheld parallel gripper. A person holds the body of the device and operates a trigger to open and close the fingers while performing a task. A GoPro mounted on top records video and data from its built-in IMU in the same MP4 file.

The main components are:

Component Role
GoPro Records RGB video together with accelerometer and gyroscope data
Fisheye lens Captures not only the area near the end effector, but also the destination and surrounding context
Side mirrors Add views from different directions within a single image, providing extra depth cues
Fiducial markers Used to estimate gripper width and align coordinate systems
Parallel gripper Lets a person demonstrate with an end effector that is closer to a robot gripper than a bare human hand
Soft fingers Provide passive compliance during contact

The reason for using a robot-like gripper instead of a bare human hand is to reduce the embodiment gap between demonstration and deployment. Because the GoPro moves together with the gripper, it is also easier to keep the observation setup consistent: both the demonstration device and the robot view the object from a wrist-level perspective.

For bimanual tasks, two identical grippers are used. Synchronizing the two GoPros makes it possible to represent the left- and right-hand motions on the same timeline.


4. What Does the UMI Camera See?

UMI images have much stronger fisheye distortion than ordinary webcam images. At first, it may seem that removing the distortion would make learning easier. UMI instead prioritizes preserving a wide field of view and uses the fisheye image directly as a policy observation.

For a cup-moving task, the camera needs to capture at least:

  • The cup currently being grasped

  • The position of the gripper tips

  • The saucer where the cup will be placed

  • The edge of the table and nearby obstacles

  • The direction the gripper should move next

With a narrow field of view, the destination can disappear from the frame as soon as the cup fills the image. A fisheye lens keeps the main object near the center while preserving more of the surrounding context in the same frame.

The image below is a 5-by-5 grid of 25 frames sampled at evenly spaced positions from camera0_rgb across the official in-the-wild cup dataset. We will cover how to download and train on this dataset later. For now, the important point is what a UMI observation actually looks like.

The backgrounds and lighting conditions vary greatly: outdoor tables, offices, wooden tables, and darker indoor spaces all appear in the dataset. The cups also vary in color, shape, and material. Even so, the overall visual structure remains fairly consistent:

  • A fisheye wrist-view image

  • The gripper area near the bottom of the frame

  • The object being manipulated near the center

  • Context about the table and destination around the edges

In other words, UMI aims to vary the environment while preserving a body-centered viewpoint that is useful to the robot. This consistency comes from carrying the same GoPro-equipped gripper from one location to another.

The Side Mirrors Are Not Stereo Cameras for SLAM

Small mirrors are visible at the lower left and right of the image. The paper describes the reflected views as implicit stereo. Because one GoPro frame contains views from additional angles, the mirrors can give the policy extra cues about depth.

However, it would not be accurate to describe UMI's SLAM system as stereo ORB-SLAM3 that treats the two mirrors as two independent cameras. The core SLAM system is a monocular-inertial ORB-SLAM3-based pipeline that uses the GoPro's fisheye video and built-in IMU. The mirrors mainly enrich the image observed by the learned policy.


5. Recovering the End-Effector Trajectory from a GoPro MP4 with SLAM

A GoPro MP4 contains not only video, but also IMU measurements such as acceleration and angular velocity. UMI uses the video and IMU data to reconstruct the 6-DoF pose of the handheld gripper over time.

Six degrees of freedom means three-dimensional position plus three-dimensional rotation:

Position: x, y, z
Rotation: three degrees of freedom corresponding to roll, pitch, and yaw

The reconstruction uses visual-inertial SLAM based on ORB-SLAM3. Combining visual features tracked in the video with IMU measurements makes scale and fast motion easier to handle than with monocular video alone.

Mapping Videos and Demonstration Videos

For each scene, UMI first records a mapping video. The gripper is moved slowly around the environment to capture the surroundings and build a map of the scene. Each demonstration video recorded afterward is then relocalized against the same map. This aligns multiple demonstrations in a shared scene coordinate system.

With the official example data, the full preprocessing sequence can be run with:

python run_slam_pipeline.py example_demo_session

This script processes the MP4 files, extracts IMU data, builds the map, runs SLAM for each demonstration, detects markers, performs calibration, and creates a plan for generating the dataset.

The main outputs are stored under example_demo_session/:

File Contents
demos/mapping/map_atlas.osa ORB-SLAM map
demos/mapping/mapping_camera_trajectory.csv Camera trajectory from the mapping video
demos/demo_*/camera_trajectory.csv Camera trajectory for each demonstration video
demos/*/imu_data.json IMU data extracted from the GoPro MP4 files
demos/*/tag_detection.pkl Fiducial-marker detection results
demos/*/slam_stdout.txt Standard output log from SLAM
demos/*/slam_stderr.txt Error log from SLAM
dataset_plan.pkl Episode, timestamp, and transform information used to generate the Zarr dataset

At this stage, the training file dataset.zarr.zip has not yet been created.

Checking SLAM Quality

The official pipeline is primarily a batch process that saves trajectory CSV files and logs, rather than automatically saving a 3D viewer visualization. To check the result, inspect the is_lost field in camera_trajectory.csv, verify that the trajectory is continuous, and look for sudden large jumps.

SLAM gives us the motion of the camera, which is also the motion of the handheld gripper. When this trajectory is combined with the gripper width estimated from marker detections, we have the time series for the demonstrated actions. The next step is to package those actions together with the images as training data.


6. Storing Images and Actions Together in Zarr

After SLAM and calibration are complete, generate the replay buffer used for training:

python scripts_slam_pipeline/07_generate_replay_buffer.py \
  -o example_demo_session/dataset.zarr.zip \
  example_demo_session

This step combines data that originally existed in separate places—MP4 files, trajectory CSV files, gripper widths, and episode boundaries—into time-aligned arrays stored in dataset.zarr.zip.

What Is Zarr?

Zarr is a format for storing large N-dimensional arrays in compressed chunks. It can be used much like a NumPy array while still allowing only the required portions to be read. This makes it well suited to machine-learning workloads that repeatedly load large image and time-series datasets.

Conceptually, a UMI Zarr dataset looks like this:

dataset.zarr.zip
├── data/
│   ├── camera0_rgb
│   ├── robot0_eef_pos
│   ├── robot0_eef_rot_axis_angle
│   ├── robot0_eef_rot_axis_angle_wrt_start
│   └── robot0_gripper_width
└── meta/
    └── episode_ends
Key Contents
camera0_rgb RGB images provided to the policy
robot0_eef_pos End-effector position
robot0_eef_rot_axis_angle End-effector rotation
robot0_eef_rot_axis_angle_wrt_start Rotation represented relative to the starting pose
robot0_gripper_width Continuous-valued gripper width
meta/episode_ends Indices that mark where each demonstration ends in the continuous arrays

The key point is that camera0_rgb[i] and the end-effector pose and gripper width at the same index represent the state at the same moment in time. This synchronization turns a video into paired observations and actions for robot imitation learning.

The JPEG XL Codec Required to Open the Zarr Dataset

UMI compresses the image arrays with JPEG XL. When opening the Zarr dataset from your own script, register the codec before reading the data:

python - <<'PY'
from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs
register_codecs()

import zarr

with zarr.ZipStore("example_demo_session/dataset.zarr.zip", mode="r") as store:
    root = zarr.open(store, mode="r")
    print(root.tree())
PY

Without registering the codec, you will see an error like this:

ValueError: codec not available: 'imagecodecs_jpegxl'

Array names alone do not make the relationship between the images and numerical values very intuitive. To make that relationship concrete, we will open one episode from the official in-the-wild dataset and follow its contents over time.


7. One Demonstration Reveals the Core Idea Behind UMI

In this section, we use episode_idx=0 from cup_in_the_wild.zarr.zip. The images, end-effector positions, and gripper widths are all extracted from the same start:end interval.

7.1 The Camera Stream Seen by the Policy

First, let us look at a video exported from camera0_rgb.

https://youtube.com/shorts/Z1tw0f1FSdE?feature=share

This video is the sequence of visual observations provided to the learning system. The object and work area appear near the center of the fisheye image, while information about the destination and table remains visible around the edges.

All of the plots below cover the same time interval as this video. Watching the video while looking at the plots makes it easier to connect the visible manipulation with the changes in the numerical data.

7.2 How the End-Effector Position Changes over Time

The plot below shows the x, y, and z components of robot0_eef_pos for the same episode.

Over roughly 40 seconds, all three position components change continuously. A section with a steep slope indicates that the end effector is moving. A flatter section may correspond to holding position while grasping the object or making a small adjustment.

The important point is that a human demonstration that originally existed only as video has now been converted into a continuous numerical trajectory corresponding to a robot end effector.

However, we should not casually interpret x, y, and z as left-right or up-down directions in the image. Their physical directions depend on the SLAM coordinate system and the coordinate transforms applied during dataset generation. These stored values are also not sent directly to the real robot as absolute-coordinate commands. During training, they can be converted into a relative representation based on the current pose, depending on the configuration.

7.3 How the Gripper Width Changes over Time

The next plot shows robot0_gripper_width over the same interval.

The width decreases sharply at around 20 seconds and increases again at around 30 seconds. Comparing the plot with the video lets us identify which parts correspond to approaching the object, grasping it, carrying it while maintaining the grip, and releasing it.

UMI records gripper width as a continuous value rather than a simple binary “open” or “closed” state. This allows intermediate openings related to object size and contact conditions to become part of the training data.

When the position trajectory and gripper width are viewed on the same timeline, the action passed to the learning system forms a recognizable sequence:

Move the end effector toward the object
        ↓
Close the gripper
        ↓
Move the end effector while maintaining the closed width
        ↓
Open the gripper at the destination

7.4 The Motion Path Viewed from Above

Projecting the x and y components onto a plane gives an intuitive top-down view of the motion across the table.

The time-series plot shows how much each axis changes, while the top-down plot shows how the end effector moves back and forth through space. Looking at the same trajectory in these two ways helps confirm that the position series recovered by SLAM is not just a collection of random values, but a coherent manipulation path.

The four pieces of information examined in this section are not separate datasets:

The same episode on the same timeline
├── camera0_rgb               : what the gripper saw
├── robot0_eef_pos            : where the end effector was
├── robot0_eef_rot_axis_angle : which way the end effector was facing
└── robot0_gripper_width      : how far the gripper was open

This time-synchronized group is the smallest building block of UMI's training data. A Diffusion Policy does not, however, process the full 40-second episode all at once. Next, we will see how short training samples are created from each episode.


8. Creating Many Training Windows from a Single Episode

In Zarr, multiple episodes are stored one after another in large continuous arrays, and meta/episode_ends marks the boundaries between them. First, let us look at the distribution of episode lengths converted to seconds at 10 Hz.

Many episodes cluster around 40 to 50 seconds, but the dataset also contains shorter demonstrations and some longer than 100 seconds. Even for the same cup-placement task, completion time changes with the environment, initial arrangement, object, and the demonstrator's operating speed.

During training, short time-series windows are extracted from these variable-length episodes by sliding the starting point through each demonstration. Because many windows can be created from one demonstration, loading the official in-the-wild dataset produced the following sample counts:

train dataset: 601274
val dataset: 33060

This does not mean that the dataset contains more than 600,000 demonstrations. The counts come from extracting many observation-history and future-action windows from roughly 1,400 variable-length demonstrations.

Windowing offers two main advantages:

  1. The model can learn from the short time range needed for the current decision instead of processing an entire long video.

  2. A single demonstration can be reused from many different starting timestamps.

At this point, the supervised data prepared by UMI is becoming concrete. Each window contains images, the current state, and the sequence of actions the human performed immediately afterward. The Diffusion Policy learns this relationship.


9. What Does a Diffusion Policy Learn?

A Diffusion Policy applies the basic idea of diffusion models used in image generation to sequences of robot actions. Noise is added to the ground-truth future actions, and the model repeatedly learns to remove that noise while conditioning on the observations. At inference time, it starts from noise and generates a plausible sequence of future actions.

Mapped to the UMI data described above, the inputs and outputs look like this:

Inputs
├── Past frames from camera0_rgb
├── Current and past end-effector poses
└── Current and past gripper widths

Outputs
├── End-effector translations and rotations for the next several steps
└── Gripper widths for the next several steps

Why Predict an Action Sequence Instead of a Single Step?

Robot manipulation is not determined by only the next instant. Picking up a cup involves a connected sequence: approach, grasp, lift, transport, and place. Predicting a short future sequence makes it easier for the model to represent this kind of coherent motion.

Why Use Relative Trajectories?

If the model memorizes an absolute command such as “move to x = 0.2 m relative to the robot base,” it becomes strongly dependent on the robot's installation position and the SLAM coordinate system. UMI instead uses a relative trajectory representation that describes which direction to move and how far, starting from the current end-effector pose.

For that reason, robot0_eef_pos as plotted in Section 7 is useful for understanding the raw data, but it is not necessarily the final action representation used by the learner. The dataset class can convert positions and rotations into relative forms, reducing dependence on the robot base frame and the world coordinate system used during data collection.

The Role of the Image Encoder

In the configuration used for this experiment, image features were extracted with vit_base_patch16_clip_224.openai, a ViT-B-family model pretrained with CLIP. Its role is to identify features relevant to manipulation even when the background, lighting, and appearance of the cup vary as much as they did in the image grid from Section 4.

We have now connected the basic ideas from the hardware all the way to the learning system. Next, we will run the full pipeline from the beginning using the official example data.


10. Running the Official Example Pipeline from MP4 to Training

The official repository provides instructions for running the system on Ubuntu 22.04.

10.1 Set Up the Environment

git clone https://github.com/real-stanford/universal_manipulation_interface.git
cd universal_manipulation_interface

sudo apt install -y libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf

mamba env create -f conda_environment.yaml
conda activate umi

10.2 Download the Example Data

wget --recursive --no-parent --no-host-directories --cut-dirs=2 \
  --relative --reject="index.html*" \
  https://real.stanford.edu/umi/data/example_demo_session/

10.3 Run SLAM and Calibration

python run_slam_pipeline.py example_demo_session

After the command finishes, confirm that files such as camera_trajectory.csv, map_atlas.osa, tag_detection.pkl, and dataset_plan.pkl have been generated.

find example_demo_session -maxdepth 3 \
  \( -name 'map_atlas.osa' \
  -o -name 'camera_trajectory.csv' \
  -o -name 'mapping_camera_trajectory.csv' \
  -o -name 'imu_data.json' \
  -o -name 'tag_detection.pkl' \
  -o -name 'dataset_plan.pkl' \) \
  -print | sort

10.4 Generate the Zarr Dataset

python scripts_slam_pipeline/07_generate_replay_buffer.py \
  -o example_demo_session/dataset.zarr.zip \
  example_demo_session

10.5 Start Training

python train.py \
  --config-name=train_diffusion_unet_timm_umi_workspace \
  task.dataset_path=example_demo_session/dataset.zarr.zip

The example session lets us verify the full connection from MP4 files through SLAM and Zarr to the learning system. Next, we will use the official in-the-wild dataset collected across many real environments and see how far we can get with a single RTX 4090.


11. Training on the Official In-the-Wild Cup Dataset

Download the official processed dataset:

wget https://real.stanford.edu/umi/data/zarr_datasets/cup_in_the_wild.zarr.zip

The image grid from Section 4, as well as the single episode and episode-duration distribution examined in Sections 7 and 8, were all read from this Zarr file. In other words, the model is trained on the same relationships among images, end-effector trajectories, and gripper widths that we have been inspecting throughout the article.

In the paper's in-the-wild cup experiment, three demonstrators collected 1,400 demonstrations across 30 environments in 12 person-hours. The paper reports 43 successful trials out of 60 in an evaluation that included unseen environments and unseen cups. The training setup was large: ViT-L/14, a total batch size of 512, 50 epochs, and eight A100 GPUs.

The differences between the paper's setup and this local experiment are summarized below.

Item In-the-wild cup setup in the paper This experiment
Data 1,400 demonstrations across 30 environments The same officially released Zarr dataset
Vision encoder CLIP ViT-L/14 vit_base_patch16_clip_224.openai
Training amount 50 epochs 1,000 training steps
Total batch size 512 4
Compute 8 x A100 1 x RTX 4090
EMA Enabled Disabled to save memory
Goal Evaluate real-robot performance Verify the data, model, and training loop

The goal here is therefore not to reproduce the paper's success rate. It is to load the official data on a local PC and confirm that the model can learn from the observation-and-action windows described in the previous sections.

11.1 The Command That Ran Reliably in This Setup

PYTHONNOUSERSITE=1 \
WANDB_MODE=disabled \
WANDB_SILENT=true \
HYDRA_FULL_ERROR=1 \
accelerate launch \
  --num_processes 1 \
  --num_machines 1 \
  --mixed_precision no \
  --dynamo_backend no \
  train.py \
  --config-name=train_diffusion_unet_timm_umi_workspace \
  task.dataset_path=cup_in_the_wild.zarr.zip \
  logging.mode=disabled \
  training.num_epochs=1 \
  training.max_train_steps=1000 \
  training.max_val_steps=2 \
  training.use_ema=False \
  dataloader.batch_size=4 \
  val_dataloader.batch_size=4 \
  dataloader.num_workers=2 \
  val_dataloader.num_workers=2 \
  dataloader.persistent_workers=True \
  val_dataloader.persistent_workers=True \
  dataloader.pin_memory=False \
  val_dataloader.pin_memory=False

The following messages appeared in the log:

Loading pretrained weights from Hugging Face hub
  (timm/vit_base_patch16_clip_224.openai)

rgb keys:
  ['camera0_rgb']

low_dim_keys keys:
  ['robot0_eef_pos',
   'robot0_eef_rot_axis_angle',
   'robot0_eef_rot_axis_angle_wrt_start',
   'robot0_gripper_width']

iterating dataset to get normalization: 100%|...| 9395/9395
train dataset: 601274 train dataloader: 150319
val dataset: 33060 val dataloader: 8265

The rgb keys and low_dim_keys entries show that the Zarr images and end-effector state fields introduced in Section 6 were recognized correctly. During iterating dataset to get normalization, the system calculates statistics used to scale positions, rotations, gripper widths, and other values into ranges that are easier for the model to learn.

train dataloader: 150319 is the number of mini-batches required to complete one full epoch with a batch size of 4. Because this run uses training.max_train_steps=1000, it stops early rather than processing all 150,319 steps.


12. What Did 1,000 Training Steps Confirm?

The plot below shows train_loss at each global step, read from the training log file logs.json.txt.

The loss started at roughly 1.0 to 1.2 and fell to around 0.1 in the later part of the run. There were occasional spikes, but the overall trend was downward.

This result confirms that the following parts of the pipeline were connected and operating together:

Read JPEG XL-compressed camera0_rgb images from Zarr
        ↓
Create mini-batches containing images and low-dimensional states from the same window
        ↓
Extract image features with a ViT
        ↓
Use the Diffusion Policy to calculate a loss for future actions
        ↓
Update the model parameters with backpropagation and an optimizer step

A lower training loss alone does not, however, prove that the real robot can successfully place cups.

What this result confirms What this result does not confirm
The learning system can read the images and states stored in Zarr Whether the policy generalizes to unseen rooms or cups
Forward and backward passes run successfully Whether the policy behaves reliably in a closed-loop real-robot rollout
The training objective decreases over 1,000 steps Whether the paper's 43/60 success result has been reproduced
A reduced setup runs on a single RTX 4090 Whether the model has converged sufficiently after 50 epochs

In real-robot imitation learning, even a small prediction error changes the next observation. That shift can then accumulate through later actions. Differences in camera lighting, contact, object slip, and latency in the robot and gripper also affect performance. Reproducing the paper's result therefore requires not only training at a comparable scale, but also safe rollout evaluation on real hardware.


13. What UMI Can Do—and Where It Has Limits

The paper applies UMI to tasks such as:

Task Challenges involved
Cup placement Grasping, pushing to change orientation, and placing a cup on a saucer
Dynamic tossing Fast motion, release timing, and latency compensation
Bimanual garment folding Coordinating two hands and manipulating deformable objects
Dishwashing Long action sequences involving water, a sponge, a faucet, and physical contact

UMI's value lies not only in the resulting policies, but also in the ability to collect task data in many locations without bringing a robot to each one. The in-the-wild cup experiment makes use of exactly the property seen in the Section 4 image grid: backgrounds, lighting, and objects can vary while the body-centered viewpoint remains consistent.

UMI also has clear limitations:

  • SLAM needs sufficient visual texture.

  • Trajectory reconstruction can become unstable around blank white walls, reflective surfaces, dark areas, or strong backlighting.

  • A demonstration performed outside the robot's reachable workspace cannot be transferred directly to the robot.

  • Collecting data with the UMI gripper is slower than demonstrating with bare human hands.

  • Real deployment requires compensation for observation and execution latency in the camera, robot, and gripper.

  • Training loss alone cannot measure closed-loop task success.

  • Reproducing the paper's results requires large-scale compute and a real-robot evaluation environment.


14. Summary

We followed UMI through the complete pipeline:

  1. A person demonstrates a task with a GoPro-equipped handheld gripper.

  2. The fisheye image becomes a wrist-view observation similar to what the robot sees during execution.

  3. SLAM reconstructs a 6-DoF trajectory from the GoPro video and IMU data.

  4. Fiducial markers are used to recover the gripper width.

  5. Images, poses, and gripper widths are placed on the same timeline and stored in Zarr.

  6. Short observation-and-action windows are extracted from variable-length episodes.

  7. A Diffusion Policy learns sequences of future relative actions.

Even without a robot, following each stage of MP4 → SLAM → Zarr → Diffusion Policy gives us a practical understanding of what UMI is doing. UMI is not merely a gripper with a camera. It is an integrated interface for turning human demonstrations collected in everyday environments into observation-and-action data that robots can use.


Appendix A. Scripts Used to Inspect the Data and Create the Plots

A.1 Image Grid Across the Full Dataset

from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs
register_codecs()

import zarr
import numpy as np
import matplotlib.pyplot as plt

path = "cup_in_the_wild.zarr.zip"

with zarr.ZipStore(path, mode="r") as store:
    root = zarr.open(store, mode="r")
    imgs = root["data"]["camera0_rgb"]
    n = imgs.shape[0]
    print("camera0_rgb:", imgs.shape, imgs.dtype)

    idxs = np.linspace(0, n - 1, 25, dtype=int)

    fig, axes = plt.subplots(5, 5, figsize=(10, 10))
    for ax, idx in zip(axes.ravel(), idxs):
        ax.imshow(imgs[idx])
        ax.set_title(f"idx={idx}", fontsize=8)
        ax.axis("off")

    plt.tight_layout()
    plt.savefig("umi_cup_wild_image_grid.png", dpi=200)

A.2 Exporting One Episode to MP4

from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs
register_codecs()

import zarr
import cv2

path = "cup_in_the_wild.zarr.zip"
episode_idx = 0
fps = 10

with zarr.ZipStore(path, mode="r") as store:
    root = zarr.open(store, mode="r")
    imgs = root["data"]["camera0_rgb"]
    episode_ends = root["meta"]["episode_ends"][:]

    start = 0 if episode_idx == 0 else int(episode_ends[episode_idx - 1])
    end = int(episode_ends[episode_idx])

    first = imgs[start]
    h, w = first.shape[:2]
    writer = cv2.VideoWriter(
        "umi_episode0_camera0.mp4",
        cv2.VideoWriter_fourcc(*"mp4v"),
        fps,
        (w, h),
    )

    for i in range(start, end):
        frame_bgr = cv2.cvtColor(imgs[i], cv2.COLOR_RGB2BGR)
        writer.write(frame_bgr)

    writer.release()

A.3 Position, Top-Down Trajectory, and Gripper Width

from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs
register_codecs()

import zarr
import numpy as np
import matplotlib.pyplot as plt

path = "cup_in_the_wild.zarr.zip"
episode_idx = 0
freq = 10.0

with zarr.ZipStore(path, mode="r") as store:
    root = zarr.open(store, mode="r")
    episode_ends = root["meta"]["episode_ends"][:]

    start = 0 if episode_idx == 0 else int(episode_ends[episode_idx - 1])
    end = int(episode_ends[episode_idx])

    pos = root["data"]["robot0_eef_pos"][start:end]
    width = root["data"]["robot0_gripper_width"][start:end].reshape(-1)

t = np.arange(len(pos)) / freq

plt.figure(figsize=(8, 4))
plt.plot(t, pos[:, 0], label="x")
plt.plot(t, pos[:, 1], label="y")
plt.plot(t, pos[:, 2], label="z")
plt.xlabel("time [s]")
plt.ylabel("position")
plt.title("End-effector position trajectory")
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("umi_episode0_eef_pos.png", dpi=200)

plt.figure(figsize=(5, 5))
plt.plot(pos[:, 0], pos[:, 1])
plt.xlabel("x")
plt.ylabel("y")
plt.title("End-effector top-down trajectory")
plt.axis("equal")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("umi_episode0_xy_trajectory.png", dpi=200)

plt.figure(figsize=(8, 3))
plt.plot(t, width)
plt.xlabel("time [s]")
plt.ylabel("gripper width")
plt.title("Gripper width trajectory")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("umi_episode0_gripper_width.png", dpi=200)

A.4 Episode-Duration Distribution

from diffusion_policy.codecs.imagecodecs_numcodecs import register_codecs
register_codecs()

import zarr
import numpy as np
import matplotlib.pyplot as plt

path = "cup_in_the_wild.zarr.zip"
freq = 10.0

with zarr.ZipStore(path, mode="r") as store:
    root = zarr.open(store, mode="r")
    episode_ends = root["meta"]["episode_ends"][:]

lengths = np.diff(np.r_[0, episode_ends])
durations = lengths / freq

print("num episodes:", len(lengths))
print("mean duration [s]:", durations.mean())
print("median duration [s]:", np.median(durations))
print("min/max duration [s]:", durations.min(), durations.max())

plt.figure(figsize=(8, 4))
plt.hist(durations, bins=50)
plt.xlabel("episode duration [s]")
plt.ylabel("count")
plt.title("Episode duration distribution")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("umi_episode_duration_hist.png", dpi=200)

A.5 Training Loss

from pathlib import Path
import json
import pandas as pd
import matplotlib.pyplot as plt

runs = sorted(
    [p for p in Path("data/outputs").glob("*/*")
     if (p / "logs.json.txt").exists()],
    key=lambda p: p.stat().st_mtime,
)

if not runs:
    raise SystemExit("Could not find a run containing logs.json.txt.")

run = runs[-1]
log_path = run / "logs.json.txt"
print("run:", run)
print("log:", log_path)

records = []
with open(log_path) as f:
    for line in f:
        line = line.strip()
        if line:
            records.append(json.loads(line))

df = pd.DataFrame(records)

plt.figure(figsize=(8, 4))
plt.plot(df["global_step"], df["train_loss"])
plt.xlabel("global step")
plt.ylabel("train loss")
plt.title("UMI training loss")
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("umi_train_loss.png", dpi=200)

References