Once demonstrations are recorded they become a dataset, and three questions decide whether it is useful: does the record contain everything needed to reproduce what happened, does it cover the states the learned policy will meet, and how will you know whether a policy trained on it works? The first is a question of format, the second of coverage, the third of evaluation, and each has a trap. A record without disturbances cannot be replayed. A larger dataset from the same distribution does not cover more of the states that matter (Lesson 14.1). And the number every training loop prints, the validation loss, measures agreement with the expert on the expert’s own states, which is not the same thing as success in closed loop. This lesson measures how different the two can be.
[!established] Dataset ecosystems Several libraries define formats for robot-learning and offline-RL datasets, among them Minari (described on PyPI as “a standard format for offline reinforcement learning datasets”, from the Farama Foundation), RLDS (“a Python library for Reinforcement Learning Datasets”) and LeRobot (Hugging Face’s library for machine learning on real-world robots). This lesson uses its own minimal format so that every field is visible; the requirements it tests apply to any of them.
This lesson has no lab; the bc lab of Lesson 14.1 shows coverage directly. Its grey demonstration paths fill only the corridors the expert drove; pushes send the clone outside them; DAgger’s blue paths fill exactly those regions. Output (2) puts numbers on that picture.
Output (1) writes 20 expert episodes recorded with pushes, one .npz file per episode plus a metadata.json that documents every array, the MuJoCo version, the timestep, a hash of the model file, a hash of the expert’s source code and the push parameters. Validation checks shapes, the action range and the terminal flags (exactly one of terminated and truncated per episode). Then each episode’s actions are replayed from its start state. With the recorded pushes, the replay reproduces every position exactly (largest error 0 m); without them, the positions drift by 0.32 m (median episode) and up to 0.57 m. A push is not an action, so a format with only states and actions loses it. The rule generalizes: every influence on the system that is not an action (disturbances, randomized parameters, resets of objects, a human nudging the scene) must be recorded, or the dataset cannot be checked against the simulator that produced it (Level 11.3’s replay check is the same idea).
Two measures of how much of the relevant state space a dataset covers:
Output (2) compares three datasets of 20 expert episodes each:
| Dataset | States | Cell coverage | Test coverage | Clone success, no pushes (3 seeds) | Clone success, pushes |
|---|---|---|---|---|---|
| starts above the obstacle only | 4 711 | 6.1% | 33.5% | 0.79, 0.95, 0.52 | 0.40, 0.39, 0.32 |
| starts everywhere | 4 722 | 12.3% | 64.8% | 1.00, 0.92, 0.99 | 0.64, 0.58, 0.59 |
| starts everywhere, with pushes | 6 840 | 21.9% | 86.5% | 1.00, 1.00, 1.00 | 0.93, 0.94, 0.92 |
The first two datasets have the same number of states, and the second covers twice as much; success under pushes follows test coverage across all three. A dataset restricted to one side of the obstacle produces clones whose success varies widely between seeds even without pushes (0.52 to 0.95), because half the test starts lie in states it never contains and what a network does there depends on its initialization. Coverage of the test distribution, which a dataset author can measure before training anything, predicted closed-loop success in this experiment better than dataset size did.
Output (3) saves clones of the last two datasets after 250, 500, 1 000, 2 000, 3 000 and 5 000 training steps, two seeds each, and scores every checkpoint twice: by its mean-squared error on 5 held-out demonstrations from the same distribution (the validation loss) and by closed-loop success with pushes on the 100 held-out starts.
On the broad dataset, the validation error falls steadily, from 0.25 to 0.05 for seed 0, while closed-loop success rises to 0.77 at 1 000 steps and then falls to 0.58 at 5 000. Fitting the demonstrations more closely made the clone worse at what demonstrations do not show. Across all 24 checkpoints the rank correlation between validation error and success is +0.43 (Spearman, p = 0.038): lower validation error went with lower success. Between datasets the disagreement is starker still: the pushed dataset’s clones have higher validation errors (it contains harder states, recoveries) and far higher success.
[!recommendation] Select and report imitation policies by closed-loop evaluation Use validation loss to detect bugs (it should fall), not to choose checkpoints or compare datasets. Choose checkpoints by closed-loop success on held-out starts under the conditions that matter (here, pushes), and report that number, with the dataset’s coverage of those conditions measured before training.
INPUT: the point-mass task and its expert; mjcourse.il.bc
PROCESS: write, validate and replay an episode dataset; measure coverage of three datasets and the success of their clones; score checkpoints by validation error and by closed-loop success
OUTPUT: a dataset directory and printed tables
```python file=examples/l14_2_offline_datasets.py “"”Lesson 14.2: demonstrations as a dataset: format, coverage, and why validation loss is not evaluation.
INPUT mjcourse.il.pointmass (task, expert), mjcourse.il.bc (cloning) PROCESS (1) write 20 expert episodes recorded with pushes in a documented episode format (one .npz per episode, metadata.json), read them back, validate them, and replay each episode’s actions from its start: with the recorded pushes and without; (2) three datasets of 20 episodes each: starts on one side of the obstacle only, starts everywhere, and starts everywhere with pushes; two coverage measures; clones trained on each (3 seeds) evaluated on 100 held-out starts with and without pushes; (3) validation loss against closed-loop success: clones of the last two datasets saved after 250 to 5 000 training steps (2 seeds), each scored by its validation error on 5 held-out demonstrations and by closed-loop success; their rank correlation OUTPUT a dataset directory and printed tables
Run: python examples/l14_2_offline_datasets.py “””
import hashlib import inspect import json import os import tempfile from multiprocessing import Pool from pathlib import Path
import mujoco import numpy as np import torch from scipy.stats import spearmanr
from mjcourse.il import bc from mjcourse.il import pointmass as pm from mjcourse.paths import model_path
FAST = os.environ.get(“MJC_FAST”) == “1” PUSH = 1.0 HELD_OUT = pm.sample_starts(np.random.default_rng(999), 30 if FAST else 100)
class RecordingPushes: “"”A random generator that remembers when pushes happened, so an episode can record them.”””
def __init__(self, seed):
self.rng, self.log, self.calls, self.pending = np.random.default_rng(seed), [], 0, None
def random(self):
self.calls += 1
return self.rng.random()
def uniform(self, low, high):
angle = self.rng.uniform(low, high)
self.log.append((self.calls - 1, angle)) # the step index at which it fired
return angle
def record_episode(start, seed, push_rate) -> dict: sim, rng = pm.PointMass(), RecordingPushes(seed) run = sim.rollout(pm.expert, start, push_rate=push_rate, rng=rng) pushes = np.array([[step, 0.4 * np.cos(a), 0.4 * np.sin(a)] for step, a in rng.log]).reshape(-1, 3) rewards = -np.linalg.norm(run.states[:, :2] - pm.GOAL, axis=1) * sim.model.opt.timestep return {“states”: run.states, “actions”: run.actions, “rewards”: rewards, “pushes”: pushes, “start”: np.asarray(start), “terminated”: np.array(run.success), “truncated”: np.array(not run.success), “seed”: np.array(seed)}
def replay(episode: dict, with_pushes: bool) -> float: “"”Largest position difference (m) between the recorded states and a replay of the actions.””” m = mujoco.MjModel.from_xml_path(str(model_path(“point_mass”))) d = mujoco.MjData(m) d.qpos[:] = episode[“start”] mujoco.mj_forward(m, d) kicks = {int(step): dv for step, *dv in episode[“pushes”]} if with_pushes else {} worst = 0.0 for t, action in enumerate(episode[“actions”]): if t in kicks: d.qvel[:] += kicks[t] worst = max(worst, float(np.abs(d.qpos - episode[“states”][t, :2]).max())) d.ctrl[:] = action mujoco.mj_step(m, d) return worst
def part1(root: Path) -> None: root.mkdir(parents=True, exist_ok=True) starts = pm.sample_starts(np.random.default_rng(3), 20) episodes = [record_episode(s, seed=100 + i, push_rate=PUSH) for i, s in enumerate(starts)] meta = {“format”: “one npz per episode”, “task”: “point mass around an obstacle (mjcourse.il.pointmass)”, “mujoco”: mujoco.version, “timestep”: 0.01, “expert_sha256”: hashlib.sha256(inspect.getsource(pm.expert).encode()).hexdigest(), “model_sha256”: hashlib.sha256(Path(model_path(“point_mass”)).read_bytes()).hexdigest(), “push”: {“rate_per_s”: PUSH, “speed_m_s”: 0.4}, “arrays”: {“states”: “(T, 4) x, y, vx, vy before each action”, “actions”: “(T, 2) force (N), in [-1, 1]”, “rewards”: “(T,) minus distance to the goal times the timestep”, “pushes”: “(P, 3) step index, dvx, dvy: disturbances that are not actions”, “start”: “(2,)”, “terminated”: “reached the goal”, “truncated”: “ran out of time”, “seed”: “push generator seed”}} (root / “metadata.json”).write_text(json.dumps(meta, indent=1)) for i, ep in enumerate(episodes): np.savez_compressed(root / f”episode_{i:03d}.npz”, *ep) loaded = [dict(np.load(p)) for p in sorted(root.glob(“episode_.npz”))] problems = [] for i, ep in enumerate(loaded): t = len(ep[“states”]) if ep[“actions”].shape != (t, 2) or ep[“rewards”].shape != (t,): problems.append(f”{i}: lengths”) if np.abs(ep[“actions”]).max() > 1.0: problems.append(f”{i}: action range”) if bool(ep[“terminated”]) == bool(ep[“truncated”]): problems.append(f”{i}: terminal flags”) with_p = [replay(ep, True) for ep in loaded] without_p = [replay(ep, False) for ep in loaded] lengths = [len(ep[“states”]) for ep in loaded] print(f” 20 episodes, {sum(lengths)} steps ({min(lengths)} to {max(lengths)} per episode), “ f”{sum(len(ep[‘pushes’]) for ep in loaded)} pushes, {sum(bool(ep[‘terminated’]) for ep in loaded)} reached the goal”) print(f” validation: {‘no problems’ if not problems else problems}”) print(f” replaying the actions from each start: largest position error with the recorded pushes {max(with_p):.1e} m; “ f”without them {np.median(without_p):.3f} m median, {max(without_p):.3f} m worst”)
def dataset(kind: str, seed: int, n: int = 20) -> tuple[np.ndarray, np.ndarray, np.ndarray]: rng = np.random.default_rng([seed, {“narrow”: 1, “broad”: 2, “broad with pushes”: 3}[kind]]) starts = pm.sample_starts(rng, n) if kind == “narrow”: starts[:, 1] = rng.uniform(0.10, 0.25, n) # one side of the obstacle only rate = PUSH if kind == “broad with pushes” else 0.0 sim = pm.PointMass() runs = [sim.rollout(pm.expert, s, push_rate=rate, rng=np.random.default_rng([seed, 9, i])) for i, s in enumerate(starts)] return np.vstack([r.states for r in runs]), np.vstack([r.actions for r in runs]), starts
def coverage(states: np.ndarray, reference: np.ndarray, scale: np.ndarray) -> tuple[float, float]: “”“(fraction of free 2 cm cells in the arena visited; fraction of reference states within 0.5 standard deviations of some dataset state).””” cells = {(int((x + 0.5) // 0.02), int((y + 0.5) // 0.02)) for x, y in states[:, :2]} xs = (np.arange(50) + 0.5) * 0.02 - 0.5 free = sum(1 for x in xs for y in xs if np.hypot(x, y) > 0.125 and abs(x) < 0.475 and abs(y) < 0.475) a, b = reference / scale, states / scale near = np.zeros(len(a), bool) for i in range(0, len(a), 1000): near[i:i + 1000] = np.sqrt(((a[i:i + 1000, None] - b[None]) ** 2).sum(-1)).min(1) < 0.5 return len(cells) / free, float(near.mean())
def train_and_score(job) -> dict: kind, seed = job torch.set_num_threads(1) states, actions, _ = dataset(kind, seed) policy, _ = bc.fit(states, actions, seed=seed) clean = pm.evaluate(policy, HELD_OUT)[“success”] pushed = pm.evaluate(policy, HELD_OUT, push_rate=PUSH)[“success”] return {“kind”: kind, “seed”: seed, “clean”: clean, “pushed”: pushed}
def part2() -> None: kinds = (“narrow”, “broad”, “broad with pushes”) reference = np.vstack([r.states for r in pm.evaluate(pm.expert, HELD_OUT, push_rate=PUSH)[“runs”]]) scale = reference.std(axis=0) print(f” {‘dataset (20 episodes)’:<22}{‘states’:>8}{‘cells covered’:>15}{‘test states covered’:>21}”) for kind in kinds: states, _, _ = dataset(kind, 0) cells, near = coverage(states, reference, scale) print(f” {kind:<22}{len(states):>8d}{cells:>15.1%}{near:>21.1%}”) jobs = [(k, s) for k in kinds for s in range(1 if FAST else 3)] if FAST: results = [train_and_score(j) for j in jobs] else: with Pool(len(jobs)) as pool: results = pool.map(train_and_score, jobs) print(“ clones, success on 100 held-out starts per seed:”) for kind in kinds: rows = [r for r in results if r[“kind”] == kind] print(f” {kind:<22} without pushes {‘, ‘.join(format(r[‘clean’], ‘.2f’) for r in rows)}; “ f”with pushes {‘, ‘.join(format(r[‘pushed’], ‘.2f’) for r in rows)}”)
def checkpoint_scores(job) -> list[dict]: kind, seed = job torch.set_num_threads(1) states, actions, _ = dataset(kind, seed) val_s, val_a, _ = dataset(kind, seed + 100, n=5) # 5 more demonstrations, same distribution rows = [] for steps in ((100, 500) if FAST else (250, 500, 1000, 2000, 3000, 5000)): policy, info = bc.fit(states, actions, steps=steps, seed=seed, validation=(val_s, val_a)) rows.append({“kind”: kind, “seed”: seed, “steps”: steps, “val”: info[“val_mse”], “success”: pm.evaluate(policy, HELD_OUT, push_rate=PUSH)[“success”]}) return rows
def part3() -> None: jobs = [(k, s) for k in (“broad”, “broad with pushes”) for s in range(1 if FAST else 2)] if FAST: rows = [r for j in jobs for r in checkpoint_scores(j)] else: with Pool(len(jobs)) as pool: rows = [r for out in pool.map(checkpoint_scores, jobs) for r in out] print(f” {‘dataset’:<20}{‘seed’:>5} validation MSE / success with pushes at 250, 500, 1000, 2000, 3000, 5000 steps”) for k, s in jobs: sel = [r for r in rows if r[“kind”] == k and r[“seed”] == s] print(f” {k:<20}{s:>5} “ + “, “.join(f”{r[‘val’]:.3f}/{r[‘success’]:.2f}” for r in sel)) rho, p = spearmanr([r[“val”] for r in rows], [r[“success”] for r in rows]) print(f” across all {len(rows)} checkpoints: Spearman correlation of validation MSE with success {rho:.2f} (p = {p:.2g})”) for k in (“broad”, “broad with pushes”): sel = [r for r in rows if r[“kind”] == k] rho_k, p_k = spearmanr([r[“val”] for r in sel], [r[“success”] for r in sel]) print(f” within ‘{k}’: {rho_k:.2f} (p = {p_k:.2g})”)
if name == “main”: print(f”MuJoCo {mujoco.version}, torch {torch.version}; pushes {PUSH} per second, 0.4 m/s”) print(“(1) an episode format, validated and replayed”) part1(Path(tempfile.gettempdir()) / “mjcourse_l14_2”) print(“(2) coverage and closed-loop success”) part2() print(“(3) validation loss against closed-loop success, per checkpoint”) part3()
Output:
```text
MuJoCo 3.14.0, torch 2.14.1+cpu; pushes 1.0 per second, 0.4 m/s
(1) an episode format, validated and replayed
20 episodes, 6149 steps (192 to 600 per episode), 50 pushes, 19 reached the goal
validation: no problems
replaying the actions from each start: largest position error with the recorded pushes 0.0e+00 m; without them 0.319 m median, 0.573 m worst
(2) coverage and closed-loop success
dataset (20 episodes) states cells covered test states covered
narrow 4711 6.1% 33.5%
broad 4722 12.3% 64.8%
broad with pushes 6840 21.9% 86.5%
clones, success on 100 held-out starts per seed:
narrow without pushes 0.79, 0.95, 0.52; with pushes 0.40, 0.39, 0.32
broad without pushes 1.00, 0.92, 0.99; with pushes 0.64, 0.58, 0.59
broad with pushes without pushes 1.00, 1.00, 1.00; with pushes 0.93, 0.94, 0.92
(3) validation loss against closed-loop success, per checkpoint
dataset seed validation MSE / success with pushes at 250, 500, 1000, 2000, 3000, 5000 steps
broad 0 0.253/0.62, 0.181/0.76, 0.145/0.77, 0.115/0.76, 0.079/0.64, 0.047/0.58
broad 1 0.340/0.63, 0.268/0.70, 0.230/0.74, 0.159/0.66, 0.114/0.58, 0.073/0.50
broad with pushes 0 0.505/0.97, 0.469/0.94, 0.412/0.96, 0.315/0.95, 0.222/0.93, 0.148/0.89
broad with pushes 1 0.377/0.94, 0.296/0.91, 0.178/0.95, 0.099/0.93, 0.100/0.94, 0.067/0.95
across all 24 checkpoints: Spearman correlation of validation MSE with success 0.43 (p = 0.038)
within 'broad': 0.36 (p = 0.24)
within 'broad with pushes': 0.41 (p = 0.19)
Replaying a dataset’s actions does not reproduce its states. Something that was not an action changed the system: disturbances, randomization, a different model or engine version. Record them, and the model hash.
A larger dataset did not help. It covers the same states. Measure coverage of the states your test conditions produce.
The checkpoint with the lowest validation loss is not the best one. Expected (output 3). Evaluate checkpoints in closed loop.
Clones vary a lot between seeds. Large regions of the test distribution have no data, and the network’s behaviour there is arbitrary. Add coverage before adding seeds.
Add a fourth dataset: 10 episodes everywhere without pushes plus 10 with pushes (similar size to the others). Predict its test coverage and the clone’s success with pushes from the table, then measure both.
Write a dataset-coverage report that a reviewer could run on any dataset in this format: per-region counts on a grid, the test coverage for a stated test distribution, the fraction of episodes with disturbances, and a list of states in the test distribution farthest from the data. Apply it to the three datasets of output (2), and to a dataset of your own from Lesson 11.3.
[!research] What validation loss misses Mandlekar et al. (arXiv:2108.03298, in the reading list) studied learning from offline human demonstrations for manipulation and list among their lessons “the variability based on the stopping criteria due to the different objectives in training and evaluation”. This lesson reproduces such a disagreement between the training objective and closed-loop success in the smallest possible setting and points to a cause: validation states come from the expert’s distribution, so the loss rewards fitting those states ever more closely while the deployed policy fails in states the validation set does not contain (inference from outputs (2) and (3)). The same argument applies to offline RL, where the value of a policy outside the data’s support is, by construction, not measured by the data.
{"id": "14.2-check", "title": "Knowledge check", "questions": [
{"kind": "mcq", "q": "A dataset stores states and actions, and replaying the actions does not reproduce the states. What is the most likely omission?",
"options": ["The reward", "Disturbances or other influences that are not actions", "The episode length", "The random seed of the network"],
"answer": 1,
"explain": "<p>Output (1): with recorded pushes the replay is exact; without them positions drift by 0.32 m.</p>"},
{"kind": "mcq", "q": "Two datasets have the same number of states. Which property predicted clone success under pushes in output (2)?",
"options": ["Number of episodes", "Coverage of the states the test conditions produce", "Validation loss", "Training time"],
"answer": 1,
"explain": "<p>33.5%, 64.8% and 86.5% test coverage against roughly 0.37, 0.60 and 0.93 success.</p>"},
{"kind": "mcq", "q": "Validation loss fell from 0.25 to 0.05 while closed-loop success fell from 0.77 to 0.58. What should select the checkpoint?",
"options": ["The lowest validation loss", "Closed-loop success on held-out starts under the test conditions", "The last checkpoint", "The training loss"],
"answer": 1,
"explain": "<p>Output (3): across 24 checkpoints, lower validation error went with lower success (Spearman +0.43).</p>"},
{"kind": "numeric", "q": "A free-space grid has 1 200 cells and a dataset's states fall in 264 of them. What is its cell coverage, in percent?",
"answer": 22, "tol": 0.1, "unit": "%",
"explain": "<p>264 / 1200 = 22%.</p>"}
]}
Level 15 replaces the state in these observations with pixels: the same cloning, from RGB and depth, and what that costs.