The datasets you will never download
Open X-Embodiment is more than a million trajectories from twenty-two different robots. DROID is seventy-six thousand demonstrations across five hundred scenes. These are the datasets real robot policies are trained on, and you cannot fit either one on a Colab T4 — OXE alone is multiple terabytes. So this chapter does not ask you to. It makes a sharper argument: every hard problem those datasets pose is already sitting in the two tiny datasets you made, and you can feel all of them — and fix one of them — offline, on a laptop, in a few minutes.
There are two problems, and they are the whole chapter:
- Cross-embodiment wrangling. OXE is not one robot. It is dozens, with different action spaces, different sensors, different formats, all poured into one training run. You have exactly this in miniature: your PushT pusher acts in 2 numbers, your ALOHA bimanual rig acts in 6. Mixing them surfaces the real questions — how do you normalize across robots, and what can a shared policy actually carry from one to another?
- Data as the bottleneck. Chapter 1.2 told you the data is the policy. At scale, that means the highest-leverage thing you can build is not a better network but a data engine — a way to turn a few demos into many valid ones. You will build a from-scratch MimicGen and measure that it works.
Cross-embodiment, made concrete
# The cross-embodiment reality, on YOUR data. Two embodiments, two datasets, two# action spaces. A real cross-embodiment mix (RT-X, OpenVLA) stacks dozens of# these; the problems are identical and already visible with two.ACT_DIM_MAX = 6 # pusht acts in 2 dims, aloha in 6 — the shared tensor is padded to the max def ensure_dataset(path: Path, gen, episodes: int, seed: int) -> None: """Generate a LeRobot demo dataset if missing (offline, deterministic), and REBUILD it if a cached one was built for a different (episodes, seed) than this run asks for. Reusing a mismatched cache would silently train on the wrong demos while metrics.json reports THIS run's config, and either field can drift: * episodes — the whole lesson is the COVERAGE-STARVED regime (a handful of source demos). A fatter cache (say ch1.1's 500-demo set) would erase it. * seed — a cache from another --seed is different data; the demos must match the seed metrics.json records. gen re-derives byte-identical demos from (episodes, seed) alone, so a rebuild is cheap and deterministic. The default path is chapter-private (see setup), so we never read or clobber a dataset another chapter wrote.""" spec = {"episodes": episodes, "seed": seed} stamp = path / "z2r_demospec.json" # the (episodes, seed) the cached demos were built for if stamp.is_file() and json.loads(stamp.read_text()) != spec: shutil.rmtree(path) # cached demos were built for a different run -> stale if not stamp.is_file(): if path.exists(): shutil.rmtree(path) # a partial/foreign dir with no stamp -> clean before regen gen.main(["--episodes", str(episodes), "--seed", str(seed), "--out", str(path), "--no-video"]) stamp.write_text(json.dumps(spec) + "\n") def load_lerobot(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Read a LeRobot v3 dataset into (obs, action, episode_index) numpy arrays — exactly the wrangling every training stack does, laid bare.""" from lerobot.datasets.lerobot_dataset import LeRobotDataset # heavy import, after cheap failures frames = LeRobotDataset("local/demos", root=path).hf_dataset.with_format("numpy") return (np.stack(frames["observation.state"]).astype(np.float32), np.stack(frames["action"]).astype(np.float32), np.asarray(frames["episode_index"])) ensure_dataset(args.pusht_data, pusht_gen_demos, args.source_episodes, args.seed)ensure_dataset(args.aloha_data, aloha_gen_demos, args.source_episodes, args.seed)pusht_obs, pusht_act, pusht_ep = load_lerobot(args.pusht_data)aloha_obs, aloha_act, aloha_ep = load_lerobot(args.aloha_data) # The heterogeneity, made numeric. Normalization stats are PER EMBODIMENT: a# pusher's 1 m/s and a gripper's open/close command do not share a scale, so one# global normalizer would crush one embodiment's signal. This is why OXE-scale# training normalizes per-dataset, and why "what transfers" is a real question —# both states are 10 numbers, but idx 2 is a block's x here and a gripper's# closedness there. A shared policy transfers STRUCTURE (a small MLP, an action# head), never the SEMANTICS of a raw dimension.embodiments = []for name, obs, act, act_dim in (("pusht", pusht_obs, pusht_act, PushTEnv.ACT_DIM), ("aloha", aloha_obs, aloha_act, 6)): embodiments.append({ "name": name, "act_dim": act_dim, "frames": int(len(obs)), "action_min": act.min(0).round(4).tolist(), "action_max": act.max(0).round(4).tolist(), }) print(f"[{name}] {len(obs)} frames, action_dim={act_dim}, " f"action range per dim {act.min(0).round(2)}..{act.max(0).round(2)}") # Zero-pad both into ONE action tensor with an action_mask marking the real dims# (the ch1.7 trick). This is the honest cost of mixing embodiments: the model# always emits 6 numbers; the mask says which ones this embodiment actually uses.n_total = len(pusht_act) + len(aloha_act)mixed_action = np.zeros((n_total, ACT_DIM_MAX), np.float32)action_mask = np.zeros((n_total, ACT_DIM_MAX), np.float32)mixed_action[:len(pusht_act), :PushTEnv.ACT_DIM] = pusht_actaction_mask[:len(pusht_act), :PushTEnv.ACT_DIM] = 1.0mixed_action[len(pusht_act):, :6] = aloha_actaction_mask[len(pusht_act):, :6] = 1.0(args.out / "cross_embodiment.json").write_text(json.dumps({ "embodiments": embodiments, "mixed_frames": n_total, "padded_action_dim": ACT_DIM_MAX, "note": "shared 10-dim state, but dims mean different things per embodiment; normalize per embodiment",}, indent=2) + "\n")print(f"mixed pile: {n_total} frames, padded action dim {ACT_DIM_MAX}, " f"mean action_mask density {action_mask.mean():.3f} (pusht wastes 4 of 6 dims)")Load your two datasets the way any training stack does — straight off disk, into plain numpy — and the heterogeneity is immediate. PushT actions live in 2 dims (pusher velocity); ALOHA actions live in 6 (two arms, two grippers). Both observations are ten numbers, and that symmetry is a trap: the layouts do not line up. Index 2 is the block's x in PushT and the right gripper's closedness in ALOHA — a length and a unitless open/close command that happen to share a slot and share nothing else. Ten-number states that mean different things, column by column, are the norm the moment you mix robots.
Two consequences fall out, and both scale straight up to OXE:
- Normalization is per embodiment. A pusher's 1 m/s and a gripper's open/close command do not share a scale. Compute one global normalizer over the mixed pile and you crush one robot's signal into the other's. Real cross-embodiment training normalizes per dataset for exactly this reason, and the per-embodiment min/max this region prints is the same statistic, by hand.
- A shared policy transfers structure, not semantics. To put both
embodiments in one action tensor you pad the narrow one up to the widest action
dim and carry an
action_maskmarking the real dims — the same zero-pad + mask you built in ch1.7. The model always emits 6 numbers; the mask tells the loss which ones a given example actually constrains, so a PushT row never trains the four ALOHA dims it left at zero. What crosses the embodiment gap is the shape of the problem (an MLP, an action head), never the meaning of raw dimension 2 — which is the honest answer to "what does a cross-embodiment model share?" and the reason OXE-scale models still need per-robot heads.
No policy is trained on the mixed pile here; that is ch1.8's job. The point is that you have now hit, on your own data, the format-wrangling and normalization reality that a single paragraph in the OXE paper glosses over.
A data engine you can actually run: MimicGen from scratch
# MimicGen-style augmentation of the PushT demos. For each source demo we read# its FIRST state, perturb the object + pusher pose, drop the env into that# perturbed start, and RE-SOLVE with the same scripted expert that made the# demos. The result is a genuinely new trajectory in the real physics — not a# fabricated one — and we keep it ONLY if the expert still succeeds. That# success filter is the honesty gate: an augmented demo the solver cannot finish# is not a demo, it is noise._JOINTS = ("tee_x", "tee_y", "tee_yaw", "pusher_x", "pusher_y") def pusht_state_obs(env: PushTEnv) -> np.ndarray: """The 10-dim PushT observation from public env props (mirrors pusht_env._obs); used to record the augmented demo's frames without touching env internals.""" px, py = env.pusher_pos tx, ty, tyaw = env.tee_pose gx, gy, gyaw = env.TARGET_POSE return np.array([px, py, tx, ty, np.sin(tyaw), np.cos(tyaw), gx, gy, np.sin(gyaw), np.cos(gyaw)], dtype=np.float32) def set_pusht_start(env: PushTEnv, seed: int, tee_xy, tee_yaw: float, pusher_xy) -> None: """Reset env (fresh counters), then place it at a chosen start via the public MuJoCo model/data + documented joint names — the MimicGen 'new object pose'.""" env.reset(seed) # resets step/success counters and mjData adr = {j: env.model.joint(j).qposadr[0] for j in _JOINTS} q = env.data.qpos q[adr["tee_x"]], q[adr["tee_y"]] = tee_xy q[adr["tee_yaw"]] = tee_yaw q[adr["pusher_x"]], q[adr["pusher_y"]] = pusher_xy env.data.ctrl[:] = 0.0 mujoco.mj_forward(env.model, env.data) def solve_from(env: PushTEnv, seed: int) -> tuple[np.ndarray, np.ndarray, bool]: """Roll the scripted expert from the env's current start; return (obs, act, success).""" expert = ScriptedExpert(noise=0.0, seed=seed) obs_list, act_list, done = [], [], False obs = pusht_state_obs(env) while not done: action = expert.action(env) obs_list.append(obs) act_list.append(action) obs, _, done, info = env.step(action) return np.asarray(obs_list, np.float32), np.asarray(act_list, np.float32), bool(info["success"]) aug_env = PushTEnv()aug_obs_all, aug_act_all, aug_ep_all = [], [], []attempts, kept = 0, 0next_ep = int(pusht_ep.max()) + 1 # augmented demos get fresh episode ids after the source demosfor src in np.unique(pusht_ep): first = pusht_obs[pusht_ep == src][0] # this source demo's initial state base_tee, base_yaw = first[2:4], float(np.arctan2(first[4], first[5])) base_pusher = first[0:2] for _ in range(args.aug_per_demo): attempts += 1 tee_xy = base_tee + rng.normal(0.0, args.aug_pos_sigma, size=2) # rescale the block onto ~the spawn annulus (a hair wider than the env's own # [0.10, 0.24] so augmentation also probes its edges): keep direction, clamp radius radius = np.hypot(*tee_xy) tee_xy = tee_xy * np.clip(radius, 0.08, 0.26) / (radius + 1e-9) tee_yaw = wrap_angle(base_yaw + rng.normal(0.0, args.aug_yaw_sigma)) pusher_xy = np.clip(base_pusher + rng.normal(0.0, args.aug_pos_sigma, size=2), -0.30, 0.30) if np.linalg.norm(pusher_xy - tee_xy) < PushTEnv._PUSHER_CLEAR: # nudge the pusher clear of the block away = (pusher_xy - tee_xy) / (np.linalg.norm(pusher_xy - tee_xy) + 1e-9) pusher_xy = np.clip(tee_xy + away * (PushTEnv._PUSHER_CLEAR + 0.02), -0.30, 0.30) set_pusht_start(aug_env, args.seed + int(src), tee_xy, tee_yaw, pusher_xy) o, a, ok = solve_from(aug_env, args.seed + int(src)) if ok: # the success filter: only physically valid, solved demos join the pile kept += 1 aug_obs_all.append(o) aug_act_all.append(a) aug_ep_all.append(np.full(len(o), next_ep)) next_ep += 1aug_yield = kept / attempts if attempts else 0.0print(f"augmentation: {kept}/{attempts} perturbed re-solves succeeded (yield {aug_yield:.2f}), " f"+{kept} demos on top of {len(np.unique(pusht_ep))} source demos")MimicGen's idea is simple and powerful: a manipulation task is object-centric, so a demonstration recorded for one object pose can be transformed into a valid demonstration for a nearby pose. Generate enough poses and a handful of human demos becomes thousands.
Our single-env analog is the most honest version of that idea available to us, because we have something MimicGen does not: a scripted expert that solves any start. So for each of your source demos, we read its initial object and pusher pose, perturb them, drop the environment into that perturbed start through the public MuJoCo model, and re-solve it with the same expert that made the originals. The trajectory that comes out is real — the true solver acting in the true physics — not a fabricated action sequence stitched onto a new start.
The critical line is the filter: we keep an augmented demo only if the expert still succeeds. A perturbed start the solver cannot finish is not a demo, it is noise, and it never joins the pile. Measured over the default config, that filter passes 96–99% of attempts — the perturbations are wide enough to add coverage, mild enough to stay solvable. That yield is itself a reading on how far you can push the object before the task stops being the task.
Does more data help? Measure it.
# Data is the policy (ch1.2), scaled. Train the SAME small BC MLP from ch1.1 on# the source demos alone, then on source + augmented, and compare rollout# success. Everything past this point is the ch1.1 recipe, deliberately.class BCPolicy(nn.Module): """3-layer MLP, obs float32[10] -> action float32[2], with normalization baked in as buffers (ch1.1). Identical for both training runs so the only variable is the DATA — the whole point of the measurement.""" def __init__(self, hidden_dim, obs_min, obs_range, act_min, act_range): super().__init__() self.net = nn.Sequential( nn.Linear(PushTEnv.OBS_DIM, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, PushTEnv.ACT_DIM), ) for name, stat in [("obs_min", obs_min), ("obs_range", obs_range), ("act_min", act_min), ("act_range", act_range)]: self.register_buffer(name, torch.from_numpy(stat)) def forward(self, obs): normalized_obs = 2.0 * (obs - self.obs_min) / self.obs_range - 1.0 normalized_action = self.net(normalized_obs.clamp(-1.0, 1.0)) return (normalized_action + 1.0) / 2.0 * self.act_range + self.act_min def train_and_eval(obs: np.ndarray, act: np.ndarray, tag: str) -> tuple[float, int]: """Fit BC on (obs, act) with the ch1.1 recipe, roll it out, return (success_rate, frames).""" # Reset the weight-init RNG so BOTH arms start from IDENTICAL weights. With the batch # order also fixed (shuffle seed below), the training DATA is the only thing that differs # between source-only and source+augmented — which is the entire point of the measurement. torch.manual_seed(args.seed) obs_min, act_min = obs.min(0), act.min(0) obs_range = np.where((obs.max(0) - obs_min) < 1e-4, np.float32(1.0), obs.max(0) - obs_min) act_range = np.where((act.max(0) - act_min) < 1e-4, np.float32(1.0), act.max(0) - act_min) policy = BCPolicy(args.hidden_dim, obs_min, obs_range, act_min, act_range).to(device) obs_t = torch.from_numpy(obs).to(device) act_t = torch.from_numpy(act).to(device) optimizer = torch.optim.Adam(policy.parameters(), lr=args.lr) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) shuffle = torch.Generator().manual_seed(args.seed) # same seed -> same batch order for both arms for epoch in range(args.epochs): for batch in torch.randperm(len(obs_t), generator=shuffle).split(args.batch_size): loss = nn.functional.mse_loss(policy(obs_t[batch]), act_t[batch]) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() env = PushTEnv() policy.eval() successes = 0 for episode in range(args.eval_episodes): obs_now, done, info = env.reset(seed=10_000 + args.seed + episode), False, {} # held-out eval seeds while not done: with torch.no_grad(): action = policy(torch.from_numpy(obs_now).to(device).unsqueeze(0))[0].cpu().numpy() obs_now, _, done, info = env.step(action) successes += bool(info["success"]) rate = successes / args.eval_episodes print(f"[{tag}] {len(obs)} frames -> eval success {successes}/{args.eval_episodes} = {rate:.2f}") return rate, len(obs) source_rate, source_frames = train_and_eval(pusht_obs, pusht_act, "source-only")if aug_obs_all: all_obs = np.concatenate([pusht_obs, *aug_obs_all]) all_act = np.concatenate([pusht_act, *aug_act_all])else: # no augmented demo survived the success filter (rare); the honest fallback is source-only all_obs, all_act = pusht_obs, pusht_actaug_rate, aug_frames = train_and_eval(all_obs, all_act, "source+augmented") if args.rerun: # the data-scale curve: success vs training-set size, both arms rr.set_time("dataset_size", sequence=source_frames) rr.log("scale/success_rate", rr.Scalars([source_rate])) rr.set_time("dataset_size", sequence=aug_frames) rr.log("scale/success_rate", rr.Scalars([aug_rate])) rr.log("scale/aug_yield", rr.Scalars([aug_yield]), static=True)This is the payoff, and it is deliberately the ch1.1 recipe, unchanged — the same
small BC MLP, normalization baked in as buffers, the same rollout eval. Two arms
are trained: once on your 12 source demos, once on those 12 plus everything the
augmentation produced. Both arms start from identical weights (the loop
re-seeds torch before each) and see batches in the same order, so the only
thing that differs between them is the data. That is what makes this a
measurement and not an anecdote — and it is why turning augmentation off
(--aug_per_demo 0) reproduces the source-only arm to the digit.
The measured result, seed sweep 0–2 at the default config:
| seed | source-only | source + augmented | delta |
|---|---|---|---|
| 0 | 0.02 | 0.30 | +0.28 |
| 1 | 0.08 | 0.20 | +0.12 |
| 2 | 0.14 | 0.26 | +0.12 |
Augmentation helps on every seed. Read the table honestly, because both halves matter:
- The ordering is the rock: source << augmented, always, +0.12 to +0.28. That is the ch1.2 thesis scaled — the policy never changed, only the data did, and more valid data bought more success.
- The absolute numbers are modest (2–30%), on purpose. Twelve demos tile almost none of the PushT spawn annulus, so source-only BC is coverage-starved; we chose that regime so the data effect is visible instead of being drowned by an already-saturated policy. This is the free-tier reality, not a state-of-the-art result, and the chapter does not pretend otherwise. With more source demos you would expect the gap to shrink as source-only stops starving — which is itself the lesson about where augmentation earns its keep: exactly where data is scarce.
One seed would have told you almost nothing here — seed 0's +0.28 next to seed 1's +0.12 — which is exactly the ch2.1 warning: read RL-flavored metrics across seeds, never off a single draw. The rerun recording logs the two success rates against training-set size — the data-scale curve, for your own eyes.
Read the real thing
Our data engine is one region of scale_data.py — augment. For each source
demo it reads the first state, perturbs the block and pusher pose, drops the env
into that start through the public MuJoCo model (set_pusht_start), and
re-solves with the same scripted expert (solve_from). The honesty gate is a
single line, if ok: — the trajectory joins the pile only when the expert still
finishes. Because we own a solver that solves any start, every kept demo is a
real solution in real physics. That ownership is also the whole simplification:
one task, one object, one embodiment, a solver on tap.
meta.yaml pins the real thing to NVlabs/mimicgen@ea09885 (tag v1.0.0). Read
it against what you just wrote.
The object-centric transform. MimicGen has no solver on tap, so instead of
re-solving it replays — geometrically. DataGenerator.generate() in
mimicgen/datagen/data_generator.py segments each source demo into per-object
subtasks, and for each one calls transform_source_data_segment_using_object_pose
in mimicgen/utils/pose_utils.py (line 261): it takes the end-effector poses that
the source demo recorded relative to the object and re-expresses them relative to
the object's new pose in the current scene. The transformed waypoints
(WaypointTrajectory) are then executed open-loop through the robot's controller.
Same object-centric bet our perturb makes — a demo for one object pose is valid
for a nearby one — but done in SE(3) pose space and run through the real
controller rather than handed back to an expert.
The success filter, at scale. Ours is if ok. Theirs is the top of
mimicgen/scripts/generate_dataset.py: after each generate(),
success = bool(generated_traj["success"]), and only on success is the episode
written and merged into the output dataset. The loop runs under guarantee_success
— keep attempting until you have collected N successful trajectories — so the
reported success_rate = num_success / num_attempts is the production twin of our
aug_yield. Same gate, industrial scale: thousands of attempts, per-episode HDF5s
merged at the end.
What they add, and why. Real tasks are multi-step — pick then place then
insert — so MimicGen segments per subtask and picks a source segment per object
frame (select_source_demo, mimicgen/datagen/selection_strategy.py); our single
push has one segment and needs none of it. And their cross-embodiment story is
robot transfer, not our normalization: transform_first_robot_pose plus the env
interface's target_pose_to_action (mimicgen/env_interfaces/base.py) let the
same object-relative waypoints drive a different arm through its own IK — the
transfer is geometric, in pose space. Be honest about the seam: the per-embodiment
min/max, zero-pad, and action_mask this chapter builds are the OXE/OpenVLA answer
to heterogeneous action vectors — a different mechanism than MimicGen's pose-space
transfer. The chapter reproduces both realities on your own tiny data; the paper
solves one of them at a million trajectories.
Read next: open mimicgen/utils/pose_utils.py and read
transform_source_data_segment_using_object_pose, then follow it up into
DataGenerator.generate() in mimicgen/datagen/data_generator.py, then out to the
success check in mimicgen/scripts/generate_dataset.py. That path — transform,
execute, keep-if-success — is exactly the three moves your augment region
compresses into a perturb, a re-solve, and an if ok.
Where this goes
You have now, offline and from scratch, hit the cross-embodiment wrangling that OXE is built on and run the MimicGen-style data engine that turns small demo sets into large ones — and you measured that it works instead of taking it on faith. The Scale Lab (meta.yaml) is this same picture at a million trajectories: stream a real OXE subset on a bigger tier and the normalization, the padding, and the "what transfers" question are identical, just larger. The read-the-real-thing segment points you at the primary sources — Open X-Embodiment, DROID, MimicGen — which you can now read as descriptions of code you have already written.