Imitation Learning Datasets: How to Structure Yours from Day One

Most robotics teams design their dataset structure twice: once implicitly, in the first week of collection, and once properly, after the first painful migration. Getting imitation learning dataset structure right on day one costs nothing extra and saves that second pass entirely.

This guide covers the five layers a dataset should have, why annotations must live outside the episode, how to name things so you can query them later, and how to design for pooling across sources you have not signed yet.

Why structure is a day-one decision

Dataset structure feels like an engineering detail you can settle later. It is not. Structure determines what you can query, what you can filter, what you can pool, and what you can explain to an auditor. Every one of those becomes urgent at a moment when changing the structure is expensive.

The pattern is predictable. A team collects for six months with a loose folder convention, then needs to answer a question like “show me every failed grasp on transparent objects under low light” and discovers the data cannot answer it. The information was present at capture and was never recorded.

The five layers of a well-structured dataset

The critical rule is the direction of dependency. Every layer is derivable from the layers above it. If you ever find yourself unable to regenerate layer four from layers one and two, something irreversible has been baked in and your dataset has become a black box.

Separate annotations from the episode

This is the structural decision that pays back the most and gets made wrong the most often.

Annotations change. Label taxonomies get revised, quality scores get recomputed, phase boundaries get corrected. If annotations live inside the episode file, every revision rewrites the episode, and you lose the ability to reproduce a model trained on an earlier labeling.

Store annotations as separate versioned artifacts that reference episodes by ID. Then a training run cites both an episode set and an annotation version, and any past result can be reproduced exactly. This is standard practice in mature annotation and labeling operations and rare in first-time ones.

Name things so future you can query them

• Use stable opaque IDs, not descriptive filenames. Descriptive names encode assumptions that become wrong. Put the description in metadata where it can be corrected.
• Record instance-level object identity. “mug_014” rather than “mug”. Costs nothing at capture, impossible to recover later.
• Use a controlled vocabulary for failure modes. Free text cannot be aggregated. Version the vocabulary and record which version applied.
• Keep one manifest per episode. A self-describing episode survives migrations that break a central index.
• Record the collection context. Rig, firmware, calibration reference, task script version, operator pseudonym, site. See demonstration log design.

Design for pooling across sources

Almost every program eventually combines sources: teleoperation from one rig, human demonstration from another, sometimes purchased data, sometimes intervention data from a deployed fleet. Pooling only works if you decided it would in advance.

1. Fix one action convention and record it explicitly per episode rather than assuming it.
2. Fix one coordinate frame convention and note the transform from each source to it.
3. Fix one time base. Absolute, monotonic, one authority.
4. Tag every episode with its source and embodiment. You will want to weight, ablate, and exclude by source later.
5. Keep native rates. Resample downstream, never at capture, or you cannot pool with a source at a different rate.

The reason this matters is that cross-source pooling is where most of the generalization gain lives, as our guide to cross-embodiment data covers.

Splits are data, not scripts

Training splits are often defined in code, regenerated each run, and never recorded. That makes results irreproducible the moment the dataset grows.

Treat a split as a first-class artifact: a named, versioned list of episode IDs, stored alongside the dataset and cited by every experiment that used it. It costs almost nothing and it is the difference between “the model got worse” and “the model got worse because the split changed.”

Split on the right axis, too. Random episode splits leak information when the same object or scene appears on both sides. Split by object instance, scene, or operator depending on what generalization you are actually claiming, then evaluate against your evaluation benchmarks.

Structure decided up front is what makes a fast ramp possible. Our mobile manipulation case study covers a 90-day path from cold-start to production.

Frequently asked questions

Should we adopt an existing dataset format?

Start from a community format for the raw and episode layers, since it makes pooling with public data far easier. Expect to extend the metadata layer, because that is where your specific operation lives.

How do we handle a mid-program schema change?

Version the schema, record the version in every episode, and write a forward migration rather than editing history. Never silently change the meaning of a field; add a new one.

What about storage cost of keeping raw streams?

Real but usually smaller than the cost of a re-collection. Tier it: fast storage during active use, cold archive afterwards. Decide retention explicitly rather than letting cost force the decision.

Who should own the dataset structure?

One named owner, typically on the ML side, with operations input. Datasets with shared ownership drift, because every team adds fields that only make sense to them.

A well-structured dataset is not more expensive to collect. It is the same data with decisions made in a different order. If you are about to start collecting and want the schema reviewed first, send us the plan.

Manish Jain · Chief Marketing Officer

Manish Jain is Chief Marketing Officer at Roborax, bringing over 20 years of experience in business strategy, digital transformation, and growth leadership to help enterprises build scalable, high-quality AI data operations.