Core concept

A dataset generated in simulation is a measurement of the simulation, and it can be wrong in the same ways a measurement can: a frame lost by the writer, a unit changed in one file, a camera pose recorded once and never updated. None of those faults is visible in a sample image, and every one of them reaches a policy trained on the data. A dataset is finished only when it comes with three things: a schema that says what every array means, in which frame and units; generation that reproduces any episode from its seed; and a validator that has been shown to fail on corrupted data.

[!established] Saving and restoring the simulation state mj_getState and mj_setState copy the parts of mjData selected by an mjtState mask. mjSTATE_INTEGRATION selects everything the next mj_step reads, including the constraint solver’s warm start (qacc_warmstart), so restoring it and applying the same controls continues the same trajectory. qpos and qvel alone are enough only where the warm start is zero, as it is right after a reset; output (2) measures what happens elsewhere.

Visual intuition

RGB, depth and labels from the front and wrist cameras at one frame of a recorded episode

Figure: one frame of episode 0, 4.8 s in, with the red cube lifted. Rows: the front and the wrist camera at 160 × 120. Columns: RGB, depth along the optical axis (metres), and the body-level label map. Regenerate with examples/l11_3_dataset.py.

Every frame of the dataset stores these six images, plus the arm’s state, the action that was held, the pose of every cube and the pose of each camera at that instant. The wrist camera moves with the hand, so its pose must be stored per frame. A single fixed pose would put the floor at the wrong height in almost every frame, which is one of the faults output (3) plants and catches.

Mathematics

The schema

Array Shape and type Meaning
time, step (T,) simulation time (s) and physics step index of each frame
qpos, qvel (T, nq), (T, nv) float64 the full state, objects included
action (T, 8) float64 joint targets (rad) and gripper half-opening command (m) held at that frame
object_pose (T, 3, 7) float64 red, green, blue cube: position (m), quaternion (w, x, y, z)
<camera>/rgb (T, H, W, 3) uint8  
<camera>/depth (T, H, W) float32 metres along the optical axis (Lesson 11.2)
<camera>/labels (T, H, W) uint8 0 background, 1 floor, 2 robot, 3 to 5 cubes, 6 tray; sites hidden
<camera>/cam_pos, cam_mat (T, 3), (T, 9) float64 cam_xpos and cam_xmat at that frame
state (T, mj_stateSize) float64 mjSTATE_INTEGRATION at each frame
physics_ctrl (steps, nu) float64 the command applied at every physics step
initial_state mj_stateSize float64 mjSTATE_INTEGRATION before the first step

A metadata.json holds what is the same for every episode: the schema version, the MuJoCo version, a SHA-256 hash of the compiled model (saved as MJB, so any change to any parameter changes it), the dataset seed and how episode seeds derive from it, the frame rate, each camera’s intrinsic matrix with its pixel convention, and the label table.

Three choices in the schema carry most of its value. The action is the command the controller tracked (joint targets and the gripper command), which is what a policy would output, while physics_ctrl stores the torques actually applied; a dataset that keeps only torques cannot train a target-level policy, and one that keeps only targets cannot be replayed. And frames are sampled at 5 Hz while the physics runs at 500 Hz, so the per-step commands are stored separately: at 8 doubles per step they cost 0.3 MB per episode, and they make exact replay possible. The per-frame integration state makes it possible to start that replay at any frame, not only at the first.

Validation

A check is worth including only if it would fail on a plausible fault. The validator in the script runs nine:

  1. the episode count against the metadata;
  2. the model hash and MuJoCo version against the current ones (a changed model or engine invalidates replay);
  3. keys, shapes and dtypes of every array;
  4. ranges: a constant frame interval, unit quaternions, finite positive depth, known labels;
  5. geometry: floor pixels back-projected with the stored depth, $K$ and camera pose must lie at $z = 0$;
  6. reprojection: each cube’s stored centre, projected with the front camera’s $K$ and pose, must lie within 2 px of its mask’s centroid (masks clear of the image border);
  7. the wrist camera: the red cube’s projected centre must land on its own label in more than 90% of the frames where it is in view;
  8. physics: the stored commands, replayed from the stored initial state, must reproduce every stored qpos and qvel bit for bit;
  9. resumption: restarting from the stored state at each episode’s middle frame must reproduce the rest of the episode bit for bit;

and, separately, regeneration: episode 0 generated again from its seed must equal the stored file array for array.

Output (2) passes all of them on 6 episodes. The cube centres lie 0.29 px from their mask centroids on average and 1.21 px at worst, the floor back-projects with a worst median error of 0.000 mm, the replay reproduces all 282 frames exactly, resuming reproduces all 138 later frames, and the regenerated episode matches in all 21 arrays, images included. The comparison line shows why the per-frame state is stored: resuming episode 0 at frame 24 from qpos and qvel alone reproduces none of the 23 later frames exactly. The largest difference is 3 × 10⁻¹¹ m, invisible in this episode, but a bitwise check can no longer tell a real fault from it.

A validator has to be tested

Output (3) corrupts a copy with three faults that real pipelines produce: frame 3 of episode 0 dropped, episode 1’s front depth written in millimetres, and episode 1’s wrist pose frozen at its first value. The schema check passes all three (every array still has a consistent shape and the right dtype), and so does the range check for the depth in millimetres, which is still finite and positive. They are caught by checks that know physics: the frame interval jumps to 0.4 s, the floor’s points land 749 m from the plane they lie on, the frozen wrist pose puts the floor at the wrong height in 9 of the 10 checked frames, and the red cube’s centre lands on its wrist mask in only 84% of the frames. Checks of form find type errors; checks of meaning find the faults that matter.

[!recommendation] Write the validator before the dataset grows Run it on every generated batch and keep a corrupted copy as its test. Check geometry (back-projection, reprojection) and physics (replay) as well as form. Record seeds, versions and a model hash with the data, and refuse to train on a dataset whose hash or version does not match its metadata.

What “reproducible” covers

The replay and the regeneration were bitwise on one machine with one build of MuJoCo and one renderer (OSMesa here). The state arrays depend on the floating-point arithmetic of the build; images also depend on the OpenGL implementation, and a different driver or GPU can change pixel values while leaving the geometry the same (inference: not tested here). Reproducibility claims should state their scope, as the metadata here does by recording versions and the model hash, and a validator that compares images across machines should compare geometry (masks, depth) with a tolerance rather than RGB bits.

Cost

Episode 0 has 48 frames and 4 700 physics steps: 15.2 MB in memory and 3.5 MB compressed. Episodes took 5.9 s each to generate here with software rendering (OSMesa, which used several CPU threads). The checkpoint’s 50 episodes therefore need about 175 MB and five minutes. Most of the size is images: the 500 Hz command stream is 0.3 MB of the 15.2, the per-frame states 23 kB.

Implementation

INPUT: Lesson 10.2's scripted pick-and-place; front and wrist cameras
PROCESS: record episodes with images, state, action, object and camera poses, and every physics command; validate form, geometry, physics and reproducibility; corrupt a copy and validate it
OUTPUT: a dataset directory, validation reports, a figure

The script imports Lesson 10.2’s pipeline from l10_2_pick_place.py in the same directory and subclasses its Episode to record while it runs.

```python file=examples/l11_3_dataset.py “"”Lesson 11.3: a pick-and-place dataset with images, and a validator that can fail it.

INPUT the scripted pick-and-place of Lesson 10.2 (imported from l10_2_pick_place.py), the front and wrist cameras of pick_place.xml PROCESS (1) generate episodes: the red cube at a seeded random pose, 5 Hz frames of RGB, depth and body-level labels from two cameras with each camera’s pose, the state, the action (joint targets and gripper command) and every object’s pose, plus the torque command of every physics step; one compressed .npz per episode and a metadata.json describing the schema, seeds, versions and model hash; (2) validate the dataset: schema and dtypes, value ranges, reprojection of cube centres into their masks, a physics replay of the stored torques that must reproduce every stored state bit for bit, and regeneration from the seed; (3) corrupt a copy in three ways and run the validator on it; (4) a figure of one frame from both cameras (runs/l11_3/frame.png) OUTPUT a dataset directory, printed reports and the figure

Run: MUJOCO_GL=egl python examples/l11_3_dataset.py [–episodes 50] [–out DIR] “””

requires: render

import argparse import hashlib import json import math import os import shutil import sys import tempfile import time from pathlib import Path

import mujoco import numpy as np

sys.path.insert(0, str(Path(file).resolve().parent)) import l10_2_pick_place as pp # noqa: E402 (the pipeline of Lesson 10.2)

from mjcourse import model_path # noqa: E402

SEED, FPS, WIDTH, HEIGHT = 11, 5, 160, 120 FIGURE = Path(file).resolve().parents[1] / “runs” / “l11_3” / “frame.png” OKABE_ITO = [“#000000”, “#999999”, “#56B4E9”, “#D55E00”, “#009E73”, “#0072B2”, “#E69F00”] # one colour per label CAMERAS = (“front”, “gripper/wrist”) LABELS = {“background”: 0, “floor”: 1, “robot”: 2, “red_cube”: 3, “green_cube”: 4, “blue_cube”: 5, “tray”: 6} SCHEMA_VERSION = “1.0”

def label_table(model: mujoco.MjModel) -> np.ndarray: “"”Geom id -> label (uint8), through each geom’s body.””” table = np.zeros(model.ngeom, np.uint8) for g in range(model.ngeom): body = model.body(model.geom_bodyid[g]).name if body in LABELS: table[g] = LABELS[body] elif model.geom(g).name == “floor”: table[g] = LABELS[“floor”] else: table[g] = LABELS[“robot”] # every other body in this scene belongs to the arm or gripper return table

def intrinsics(model: mujoco.MjModel, camera: str) -> np.ndarray: f = 0.5 * HEIGHT / math.tan(math.radians(model.cam_fovy[model.camera(camera).id]) / 2) return np.array([[f, 0, WIDTH / 2], [0, f, HEIGHT / 2], [0, 0, 1]])

class RecordingEpisode(pp.Episode): “"”Lesson 10.2’s Episode, recording frames at FPS and the torque command of every step.”””

def __init__(self, cube_xy: np.ndarray, yaw: float):
    super().__init__(cube_xy, yaw)
    m = self.model
    self.renderer = mujoco.Renderer(m, HEIGHT, WIDTH)
    self.option = mujoco.MjvOption()
    self.option.sitegroup[:] = 0                                   # sites are not objects (Lesson 11.2)
    self.labels = label_table(m)
    self.every = round(1.0 / (FPS * m.opt.timestep))
    self.steps = 0
    spec = mujoco.mjtState.mjSTATE_INTEGRATION                     # qpos, qvel, act, ctrl, warmstart, time, ...
    self.initial_state = np.zeros(mujoco.mj_stateSize(m, spec))
    mujoco.mj_getState(m, self.data, self.initial_state, spec)
    self.physics_ctrl: list[np.ndarray] = []
    self.frames: dict[str, list] = {}
    self.record()

def step(self, n: int = 1) -> None:
    for _ in range(n):
        super().step(1)
        self.physics_ctrl.append(self.data.ctrl.copy())
        self.steps += 1
        if self.steps % self.every == 0:
            self.record()

def add(self, key: str, value) -> None:
    self.frames.setdefault(key, []).append(np.asarray(value).copy())

def record(self) -> None:
    m, d = self.model, self.data
    self.add("time", d.time)
    self.add("step", self.steps)
    self.add("qpos", d.qpos)
    self.add("qvel", d.qvel)
    state = np.zeros(mujoco.mj_stateSize(m, mujoco.mjtState.mjSTATE_INTEGRATION))
    mujoco.mj_getState(m, d, state, mujoco.mjtState.mjSTATE_INTEGRATION)
    self.add("state", state)
    self.add("action", np.r_[self.target, d.ctrl[7]])               # joint targets (7) + gripper command
    self.add("object_pose", np.stack([np.r_[d.xpos[b], d.xquat[b]] for b in
                                      (m.body("red_cube").id, m.body("green_cube").id, m.body("blue_cube").id)]))
    for cam in CAMERAS:
        key = cam.split("/")[-1]
        c = m.camera(cam).id
        self.add(f"{key}/cam_pos", d.cam_xpos[c])
        self.add(f"{key}/cam_mat", d.cam_xmat[c])
        r = self.renderer
        r.disable_depth_rendering()
        r.disable_segmentation_rendering()
        r.update_scene(d, cam, scene_option=self.option)
        self.add(f"{key}/rgb", r.render())
        r.enable_depth_rendering()
        r.update_scene(d, cam, scene_option=self.option)
        self.add(f"{key}/depth", r.render())
        r.disable_depth_rendering()
        r.enable_segmentation_rendering()
        r.update_scene(d, cam, scene_option=self.option)
        seg = r.render()
        labels = np.zeros(seg.shape[:2], np.uint8)
        hit = (seg[..., 1] == mujoco.mjtObj.mjOBJ_GEOM) & (seg[..., 0] >= 0)
        labels[hit] = self.labels[seg[..., 0][hit]]
        self.add(f"{key}/labels", labels)

def model_hash() -> str: “"”sha256 of the compiled model, saved as MJB: changes whenever anything in the model does.””” with tempfile.TemporaryDirectory() as tmp: path = os.path.join(tmp, “m.mjb”) mujoco.mj_saveModel(mujoco.MjModel.from_xml_path(str(model_path(“pick_place”))), path, None) return hashlib.sha256(Path(path).read_bytes()).hexdigest()

def episode_state(seed: int, index: int) -> tuple[np.ndarray, float]: rng = np.random.default_rng([seed, index]) return np.array([rng.uniform(0.40, 0.60), rng.uniform(-0.20, 0.00)]), float(rng.uniform(-math.pi / 4, math.pi / 4))

def generate_episode(seed: int, index: int) -> dict[str, np.ndarray]: cube_xy, yaw = episode_state(seed, index) ep = RecordingEpisode(cube_xy, yaw) outcome, attempts, tipped = pp.run_episode(cube_xy, yaw, episode=ep) ep.renderer.close() arrays = {k: np.stack(v) for k, v in ep.frames.items()} arrays[“initial_state”] = ep.initial_state arrays[“physics_ctrl”] = np.stack(ep.physics_ctrl) arrays[“outcome”] = np.array(outcome) arrays[“cube_init”] = np.r_[cube_xy, yaw] return arrays

def write_dataset(out: Path, episodes: int, seed: int = SEED) -> dict: out.mkdir(parents=True, exist_ok=True) model = mujoco.MjModel.from_xml_path(str(model_path(“pick_place”))) meta = { “schema_version”: SCHEMA_VERSION, “mujoco_version”: mujoco.version, “model”: “pick_place.xml”, “model_sha256”: model_hash(), “seed”: seed, “episode_seeds”: “numpy.random.default_rng([seed, index])”, “episodes”: episodes, “fps”: FPS, “timestep”: model.opt.timestep, “image_size”: [WIDTH, HEIGHT], “cameras”: {cam.split(“/”)[-1]: {“name”: cam, “K”: intrinsics(model, cam).tolist(), “convention”: “pixel centres at index + 0.5; cam_mat is MuJoCo’s cam_xmat (camera looks along -z, y up)”} for cam in CAMERAS}, “labels”: LABELS, “depth”: “float32 metres along the optical axis; sites hidden in RGB, depth and labels”, “arrays”: {“time”: “(T,) s”, “step”: “(T,) physics step index”, “qpos”: “(T, nq)”, “qvel”: “(T, nv)”, “state”: “(T, mj_stateSize) mj_getState(mjSTATE_INTEGRATION) at each frame”, “action”: “(T, 8) joint targets (rad) and gripper half-opening command (m), as held at that frame”, “object_pose”: “(T, 3, 7) red, green, blue cube: position (m) and quaternion (w, x, y, z)”, “/rgb": "(T, H, W, 3) uint8", "/depth": "(T, H, W) float32 m", "/labels": "(T, H, W) uint8", "/cam_pos": "(T, 3)", "/cam_mat": "(T, 9)", "physics_ctrl": "(steps, nu) torque and gripper command applied at every physics step", "initial_state": "mj_getState(mjSTATE_INTEGRATION) before the first step", "outcome": "scripted pipeline's outcome", "cube_init": "(x, y, yaw)"}, "generator": "examples/l11_3_dataset.py", } (out / "metadata.json").write_text(json.dumps(meta, indent=1)) for i in range(episodes): np.savez_compressed(out / f"episode_{i:04d}.npz", **generate_episode(seed, i)) return meta

—————————————————————- validation

def project(k: np.ndarray, cam_pos: np.ndarray, cam_mat: np.ndarray, p: np.ndarray) -> tuple[np.ndarray, float]: r_wc = cam_mat.reshape(3, 3) @ np.diag([1.0, -1.0, -1.0]) pc = r_wc.T @ (p - cam_pos) return (k @ (pc / pc[2]))[:2], pc[2]

def floor_height(k, cam_pos, cam_mat, depth, labels) -> np.ndarray: “"”Heights of the floor pixels’ back-projected points: zero if depth, K and pose agree.””” vs, us = np.nonzero(labels == LABELS[“floor”]) z = depth[vs, us] p_cv = np.c_[(us + 0.5 - k[0, 2]) * z / k[0, 0], (vs + 0.5 - k[1, 2]) * z / k[1, 1], z] r_wc = cam_mat.reshape(3, 3) @ np.diag([1.0, -1.0, -1.0]) return (p_cv @ r_wc.T + cam_pos)[:, 2]

def validate(root: Path, regenerate: int | None = 0) -> list[tuple[str, bool, str]]: “"”Run every check; returns (check, passed, detail) rows.””” meta = json.loads((root / “metadata.json”).read_text()) files = sorted(root.glob(“episode_*.npz”)) model = mujoco.MjModel.from_xml_path(str(model_path(“pick_place”))) w, h = meta[“image_size”] fail: dict[str, list[str]] = {k: [] for k in (“schema”, “ranges”, “floor”, “replay”, “resume”)} reproj, inside, floor_err, replay_frames, resume_frames = [], [], [], 0, 0 for path in files: z = np.load(path) t = len(z[“time”]) n_state = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_INTEGRATION) expected = {“qpos”: ((t, model.nq), np.float64), “qvel”: ((t, model.nv), np.float64), “action”: ((t, 8), np.float64), “state”: ((t, n_state), np.float64), “object_pose”: ((t, 3, 7), np.float64)} for cam in meta[“cameras”]: expected |= {f”{cam}/rgb”: ((t, h, w, 3), np.uint8), f”{cam}/depth”: ((t, h, w), np.float32), f”{cam}/labels”: ((t, h, w), np.uint8), f”{cam}/cam_pos”: ((t, 3), np.float64), f”{cam}/cam_mat”: ((t, 9), np.float64)} for key, (shape, dtype) in expected.items(): if key not in z.files or z[key].shape != shape or z[key].dtype != dtype: fail[“schema”].append(f”{path.name}:{key}”) if not np.allclose(np.diff(z[“time”]), 1.0 / meta[“fps”], atol=1e-9): fail[“ranges”].append(f”{path.name}: frame interval {np.diff(z[‘time’]).max():.2f} s”) if not np.allclose(np.linalg.norm(z[“object_pose”][…, 3:], axis=-1), 1.0, atol=1e-9): fail[“ranges”].append(f”{path.name}: quaternion norm”) for cam, spec in meta[“cameras”].items(): depth, labels, k = z[f”{cam}/depth”], z[f”{cam}/labels”], np.array(spec[“K”]) if not (np.isfinite(depth).all() and depth.min() > 0) or labels.max() > max(meta[“labels”].values()): fail[“ranges”].append(f”{path.name}:{cam} depth or labels”) for f in range(0, t, 5): # every fifth frame is enough to catch a fault heights = np.abs(floor_height(k, z[f”{cam}/cam_pos”][f], z[f”{cam}/cam_mat”][f], depth[f], labels[f])) if len(heights): floor_err.append(np.median(heights)) if np.median(heights) > 0.001: fail[“floor”].append(f”{path.name}:{cam} frame {f}: {1000 * np.median(heights):.0f} mm”) k = np.array(meta[“cameras”][“front”][“K”]) for f in range(t): # cube centres against their front masks for j, name in enumerate((“red_cube”, “green_cube”, “blue_cube”)): vs, us = np.nonzero(z[“front/labels”][f] == meta[“labels”][name]) if len(us) < 20 or us.min() == 0 or vs.min() == 0 or us.max() == w - 1 or vs.max() == h - 1: continue uv, _ = project(k, z[“front/cam_pos”][f], z[“front/cam_mat”][f], z[“object_pose”][f, j, :3]) reproj.append(np.hypot(uv[0] - us.mean() - 0.5, uv[1] - vs.mean() - 0.5)) kw = np.array(meta[“cameras”][“wrist”][“K”]) for f in range(t): # the red cube’s centre lands on its wrist mask uv, depth_c = project(kw, z[“wrist/cam_pos”][f], z[“wrist/cam_mat”][f], z[“object_pose”][f, 0, :3]) u, v = int(uv[0]), int(uv[1]) if depth_c > 0 and 0 <= u < w and 0 <= v < h: inside.append(z[“wrist/labels”][f, v, u] == meta[“labels”][“red_cube”]) data = mujoco.MjData(model) # replay the stored torques mujoco.mj_setState(model, data, z[“initial_state”], mujoco.mjtState.mjSTATE_INTEGRATION) frame_at = {int(s): i for i, s in enumerate(z[“step”])} if 0 in frame_at and not np.array_equal(data.qpos, z[“qpos”][frame_at[0]]): fail[“replay”].append(f”{path.name}: initial state”) for s, ctrl in enumerate(z[“physics_ctrl”], start=1): data.ctrl[:] = ctrl mujoco.mj_step(model, data) if s in frame_at: replay_frames += 1 if not (np.array_equal(data.qpos, z[“qpos”][frame_at[s]]) and np.array_equal(data.qvel, z[“qvel”][frame_at[s]])): fail[“replay”].append(f”{path.name}: step {s}”) mid = t // 2 # resuming from a stored frame is exact too same, _ = resume(model, z, mid, full_state=True) resume_frames += t - mid - 1 if same != t - mid - 1: fail[“resume”].append(f”{path.name}: {same}/{t - mid - 1}”) reproj = np.array(reproj) rows = [(“episode count matches metadata”, len(files) == meta[“episodes”], f”{len(files)} files, {meta[‘episodes’]} declared”), (“model and MuJoCo version match”, model_hash() == meta[“model_sha256”] and mujoco.version == meta[“mujoco_version”], f”{meta[‘model_sha256’][:12]}, {meta[‘mujoco_version’]}”), (“schema: keys, shapes, dtypes”, not fail[“schema”], “; “.join(fail[“schema”][:3]) or “all arrays”), (“ranges: frame timing, quaternions, depth, labels”, not fail[“ranges”], “; “.join(fail[“ranges”][:3]) or “all in range”), (“floor pixels back-project to z = 0 (median < 1 mm)”, not fail[“floor”], “, “.join(f”{sum(f’:{cam} ‘ in n for n in fail[‘floor’])} {cam} frames fail” for cam in meta[“cameras”]) if fail[“floor”] else f”{len(floor_err)} frame-cameras, worst median {1000 * max(floor_err):.3f} mm”), (“front: cube centres within 2 px of mask centroids”, bool(len(reproj)) and reproj.max() < 2.0, f”{len(reproj)} cube-frames, mean {reproj.mean():.2f} px, max {reproj.max():.2f} px” if len(reproj) else “no complete masks”), (“wrist: red cube centre lands on its mask (> 90%)”, bool(inside) and np.mean(inside) > 0.9, f”{np.mean(inside):.1%} of {len(inside)} frames with the cube in view” if inside else “never in view”), (“physics replay reproduces states bit for bit”, not fail[“replay”] and replay_frames > 0, “; “.join(fail[“replay”][:2]) or f”{replay_frames} frames identical”), (“resuming at each episode’s middle frame is exact”, not fail[“resume”] and resume_frames > 0, “; “.join(fail[“resume”][:2]) or f”{resume_frames} later frames identical”)] if regenerate is not None: fresh = generate_episode(meta[“seed”], regenerate) stored = np.load(files[regenerate]) differing = [k for k in stored.files if not np.array_equal(stored[k], fresh[k])] rows.append((f”episode {regenerate} regenerated from its seed is identical”, not differing, f”{len(stored.files)} arrays compared” + (f”; differ: {differing[:4]}” if differing else “”))) return rows

def resume(model, z, frame: int, full_state: bool) -> tuple[int, float]: “"”Restart at a stored frame and step to the end with the stored commands; returns the number of later frames reproduced exactly and the largest qpos difference.””” data = mujoco.MjData(model) if full_state: mujoco.mj_setState(model, data, z[“state”][frame], mujoco.mjtState.mjSTATE_INTEGRATION) else: start = int(z[“step”][frame]) data.qpos[:], data.qvel[:], data.time = z[“qpos”][frame], z[“qvel”][frame], float(z[“time”][frame]) data.ctrl[:] = z[“physics_ctrl”][start - 1] if start > 0 else 0.0 frame_at = {int(s): i for i, s in enumerate(z[“step”])} same, worst = 0, 0.0 for s in range(int(z[“step”][frame]) + 1, len(z[“physics_ctrl”]) + 1): data.ctrl[:] = z[“physics_ctrl”][s - 1] mujoco.mj_step(model, data) if s in frame_at: i = frame_at[s] same += bool(np.array_equal(data.qpos, z[“qpos”][i]) and np.array_equal(data.qvel, z[“qvel”][i])) worst = max(worst, float(np.abs(data.qpos - z[“qpos”][i]).max())) return same, worst

def show(rows) -> None: for name, ok, detail in rows: print(f” [{‘pass’ if ok else ‘FAIL’}] {name:<54} {detail}”)

def corrupt(src: Path, dst: Path) -> None: “"”Three realistic faults: a dropped frame, depth in millimetres, and a stale camera pose.””” shutil.copytree(src, dst) files = sorted(dst.glob(“episode_*.npz”)) z = dict(np.load(files[0])) for key in list(z): if key in (“time”, “step”, “qpos”, “qvel”, “state”, “action”, “object_pose”) or “/” in key: z[key] = np.delete(z[key], 3, axis=0) # frame 3 lost by a writer np.savez_compressed(files[0], **z) z = dict(np.load(files[1])) z[“front/depth”] = (z[“front/depth”] * 1000).astype(np.float32) # one episode written in mm z[“wrist/cam_pos”] = np.repeat(z[“wrist/cam_pos”][:1], len(z[“wrist/cam_pos”]), axis=0) # pose not updated z[“wrist/cam_mat”] = np.repeat(z[“wrist/cam_mat”][:1], len(z[“wrist/cam_mat”]), axis=0) np.savez_compressed(files[1], **z)

def figure(root: Path, out: Path) -> int: “"”RGB, depth and labels from both cameras at the frame where the red cube is highest.””” import matplotlib matplotlib.use(“Agg”) import matplotlib.pyplot as plt from matplotlib.colors import ListedColormap

z = np.load(sorted(root.glob("episode_*.npz"))[0])
f = int(np.argmax(z["object_pose"][:, 0, 2]))
fig, axes = plt.subplots(2, 3, figsize=(10, 5.4), constrained_layout=True)
for row, cam in enumerate(("front", "wrist")):
    axes[row, 0].imshow(z[f"{cam}/rgb"][f])
    depth = np.ma.masked_where(z[f"{cam}/labels"][f] == LABELS["background"], z[f"{cam}/depth"][f])
    im = axes[row, 1].imshow(depth, cmap="viridis")
    fig.colorbar(im, ax=axes[row, 1], fraction=0.046, label="depth (m)")
    axes[row, 2].imshow(z[f"{cam}/labels"][f], cmap=ListedColormap(OKABE_ITO), vmin=-0.5, vmax=len(LABELS) - 0.5, interpolation="nearest")
    axes[row, 0].set_ylabel(cam)
for col, title in enumerate(("RGB", "depth along the optical axis", "labels")):
    axes[0, col].set_title(title)
for ax in axes.ravel():
    ax.set_xticks([])
    ax.set_yticks([])
handles = [plt.Rectangle((0, 0), 1, 1, color=OKABE_ITO[v]) for v in LABELS.values()]
fig.legend(handles, list(LABELS), loc="lower center", ncol=len(LABELS), frameon=False, bbox_to_anchor=(0.5, -0.06))
out.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(out, dpi=150, bbox_inches="tight")
plt.close(fig)
return f

if name == “main”: ap = argparse.ArgumentParser() ap.add_argument(“–episodes”, type=int, default=2 if os.environ.get(“MJC_FAST”) == “1” else 6) ap.add_argument(“–out”, type=Path, default=Path(tempfile.gettempdir()) / “mjcourse_l113”) args = ap.parse_args() shutil.rmtree(args.out, ignore_errors=True) print(f”MuJoCo {mujoco.version}; {args.episodes} episodes, seed {SEED}, {FPS} Hz, {WIDTH} x {HEIGHT}, cameras {CAMERAS}”) print(“(1) generate”) start = time.perf_counter() write_dataset(args.out, args.episodes) elapsed = time.perf_counter() - start sizes = [p.stat().st_size for p in sorted(args.out.glob(“episode.npz”))] first = np.load(sorted(args.out.glob(“episode_.npz”))[0]) raw = sum(first[k].nbytes for k in first.files) outcomes = [str(np.load(p)[“outcome”]) for p in sorted(args.out.glob(“episode_.npz”))] print(f” {elapsed:.1f} s for {args.episodes} episodes ({elapsed / args.episodes:.1f} s each); outcomes: {outcomes}”) print(f” episode 0: {len(first[‘time’])} frames, {len(first[‘physics_ctrl’])} physics steps; “ f”{raw / 1e6:.1f} MB in memory, {sizes[0] / 1e6:.1f} MB compressed; all episodes {sum(sizes) / 1e6:.1f} MB”) for key in (“front/rgb”, “front/depth”, “front/labels”, “wrist/rgb”, “physics_ctrl”): print(f” {key:<14} {str(first[key].shape):<22} {first[key].dtype}”) print(“(2) validate”) show(validate(args.out)) ep0 = np.load(sorted(args.out.glob(“episode_.npz”))[0]) mid = len(ep0[“time”]) // 2 same, worst = resume(mujoco.MjModel.from_xml_path(str(model_path(“pick_place”))), ep0, mid, full_state=False) print(f” for comparison, resuming episode 0 at frame {mid} from qpos and qvel alone: {same} of “ f”{len(ep0[‘time’]) - mid - 1} later frames identical, largest qpos difference {worst:.1e}”) print(“(3) a corrupted copy: frame 3 dropped in episode 0; episode 1 with front depth in mm and a frozen wrist pose”) bad = args.out.with_name(args.out.name + “_corrupted”) shutil.rmtree(bad, ignore_errors=True) corrupt(args.out, bad) show(validate(bad, regenerate=None)) frame = figure(args.out, FIGURE) print(f”(4) figure of episode 0, frame {frame} (time {frame / FPS:.1f} s): {FIGURE.relative_to(FIGURE.parents[2])}”)

Output:

```text
MuJoCo 3.14.0; 6 episodes, seed 11, 5 Hz, 160 x 120, cameras ('front', 'gripper/wrist')
(1) generate
  35.2 s for 6 episodes (5.9 s each); outcomes: ['success', 'success', 'success', 'success', 'success', 'success']
  episode 0: 48 frames, 4700 physics steps; 15.2 MB in memory, 3.5 MB compressed; all episodes 21.5 MB
    front/rgb      (48, 120, 160, 3)      uint8
    front/depth    (48, 120, 160)         float32
    front/labels   (48, 120, 160)         uint8
    wrist/rgb      (48, 120, 160, 3)      uint8
    physics_ctrl   (4700, 8)              float64
(2) validate
  [pass] episode count matches metadata                         6 files, 6 declared
  [pass] model and MuJoCo version match                         70807f4d5b78, 3.14.0
  [pass] schema: keys, shapes, dtypes                           all arrays
  [pass] ranges: frame timing, quaternions, depth, labels       all in range
  [pass] floor pixels back-project to z = 0 (median < 1 mm)     120 frame-cameras, worst median 0.000 mm
  [pass] front: cube centres within 2 px of mask centroids      288 cube-frames, mean 0.29 px, max 1.21 px
  [pass] wrist: red cube centre lands on its mask (> 90%)       100.0% of 288 frames with the cube in view
  [pass] physics replay reproduces states bit for bit           282 frames identical
  [pass] resuming at each episode's middle frame is exact       138 later frames identical
  [pass] episode 0 regenerated from its seed is identical       21 arrays compared
  for comparison, resuming episode 0 at frame 24 from qpos and qvel alone: 0 of 23 later frames identical, largest qpos difference 3.1e-11
(3) a corrupted copy: frame 3 dropped in episode 0; episode 1 with front depth in mm and a frozen wrist pose
  [pass] episode count matches metadata                         6 files, 6 declared
  [pass] model and MuJoCo version match                         70807f4d5b78, 3.14.0
  [pass] schema: keys, shapes, dtypes                           all arrays
  [FAIL] ranges: frame timing, quaternions, depth, labels       episode_0000.npz: frame interval 0.40 s
  [FAIL] floor pixels back-project to z = 0 (median < 1 mm)     10 front frames fail, 9 wrist frames fail
  [pass] front: cube centres within 2 px of mask centroids      287 cube-frames, mean 0.29 px, max 1.21 px
  [FAIL] wrist: red cube centre lands on its mask (> 90%)       84.3% of 287 frames with the cube in view
  [pass] physics replay reproduces states bit for bit           281 frames identical
  [pass] resuming at each episode's middle frame is exact       138 later frames identical
(4) figure of episode 0, frame 24 (time 4.8 s): runs/l11_3/frame.png

Debugging

A policy trains on the dataset but behaves as if the camera were somewhere else. Per-frame camera poses missing or frozen for a moving camera. The floor back-projection check finds it.

Replay or resumption does not reproduce the stored states exactly. The state was saved as qpos and qvel only, which is enough right after a reset but not mid-episode, where the solver’s warm start is nonzero. Save mjSTATE_INTEGRATION, and store the command of every physics step, not only of the recorded frames.

Regenerating an episode gives different numbers. Randomness drawn from a global generator, or from one generator shared across episodes so that episode $i$ depends on episodes before it. Derive each episode’s generator from the dataset seed and the episode index.

Depth values that are plausible but wrong. A unit or a scale applied to some files and not others. Range checks will not see it; the floor will.

Labels that change meaning between datasets. The label table lives in code and was edited. Store it in the metadata and check it.

Exercise

Generate the checkpoint’s dataset: python examples/l11_3_dataset.py --episodes 50 --out runs/l11_3_full. Time it, check the size against the estimate above, and add a check of your own that would fail if the action column were shifted by one frame.

Challenge

Make the dataset loadable by a learning library without copying it: write a PyTorch Dataset (or a tf.data pipeline) that reads frames lazily from the .npz files, returns (images, state, action) tuples with the action of the following frame as the target, and passes this lesson’s validator on its own output. Measure frames per second read, with and without decompression caching.

Research connection

[!research] Datasets as results Robot-learning datasets are reused far beyond the work that produced them, and a fault in one propagates into every policy trained on it. Large multi-robot collections such as Open X-Embodiment (Lesson 16’s reading) merge datasets with different camera conventions, action spaces and frame rates, and the quality of those conversions is rarely measured. The checks of this lesson (geometry, replay, reproducibility) apply directly to simulated subsets, and the geometric ones apply to real data whenever calibration and object poses are recorded (inference). A dataset release that ships its validator, and the corrupted copies it was tested on, lets others verify it rather than trust it.

{"id": "11.3-check", "title": "Knowledge check", "questions": [
  {"kind": "mcq", "q": "You restart a recorded episode from a mid-episode frame, setting qpos and qvel from the file and applying the stored commands, and no later frame matches the file bit for bit. What is missing?",
   "options": ["The random seed", "The rest of the integration state, such as the constraint solver's warm start", "The camera images", "The model's timestep"],
   "answer": 1,
   "explain": "<p>mjSTATE_INTEGRATION includes qacc_warmstart and everything else mj_step reads. Output (2): 0 of 23 frames match from qpos and qvel alone, all 138 from the stored state.</p>"},
  {"kind": "mcq", "q": "Which check caught depth written in millimetres in output (3)?",
   "options": ["The schema check", "The range check", "Back-projecting floor pixels to z = 0", "The model hash"],
   "answer": 2,
   "explain": "<p>Millimetre depth is still finite, positive float32 of the right shape; only geometry reveals it.</p>"},
  {"kind": "numeric", "q": "Frames are recorded at 5 Hz from a 500 Hz simulation. How many physics steps lie between consecutive frames?",
   "answer": 100, "tol": 0, "unit": "steps",
   "explain": "<p>500 / 5 = 100: why the per-step commands must be stored separately for an exact replay.</p>"},
  {"kind": "numeric", "q": "One episode compresses to 3.5 MB. About how many megabytes do 50 episodes take?",
   "answer": 175, "tol": 10, "unit": "MB",
   "explain": "<p>50 × 3.5 = 175 MB; episodes differ slightly in length, so the total varies a little.</p>"}
]}

Next

Level 12 wraps the pick-and-place and push tasks as Gymnasium environments: observations, actions, rewards and terminations, and a test that no observation reads simulator state a real robot could not measure.