See it work
That blue T-shape is being pushed onto the green target by a neural network. Nobody told it what a "T" is, or where the target is, or which way to shove. It was shown a few hundred examples of a pushing task being done, and asked to copy them. And this is the copy, running on a start it was never shown. No planner, no reward, no search. A stack of matrix multiplies that learned to push.
You are going to train that network yourself, right now, before we explain a single thing about how it works. That order is deliberate. The rest of Phase 0 takes the simulator, the data format, and the learning apart piece by piece; it reads very differently once you've already watched the whole thing succeed.
The whole loop, once
Everything is in one file, quickstart.py, and it does four things in order:
gets demonstrations, trains a network to imitate them, rolls the result out on
fresh starts, and prints how often it worked. Run it:
python curriculum/phase0_foundations/ch0.0_quickstart/quickstart.py --seed 0
| cpu | expected wall-clock on cpu-laptop: ~2.83 min (measured) | measured |
| mps | expected wall-clock on mps: ~2.84 min (measured) | measured |
| t4 | expected wall-clock on t4: ~1.7 min (measured) | measured |
| l40s | expected wall-clock on l40s: ~0.94 min (measured) | measured |
Demonstrations. A demonstration is just a recording of the task being done well: at each instant, what the world looked like and what the driver did about it. Here the "driver" is a scripted expert that ships with the PushT environment, a hand-tuned controller somebody wrote so you don't have to. We run it a few hundred times and keep what it saw and did.
# A demonstration is a list of (observation, action) pairs — what the world# looked like, and what a competent driver did about it. In ch0.4 that driver# is YOU, with a mouse. Here it's the scripted expert from the shared PushT env# (a hand-tuned controller that took real engineering to write — that effort is# exactly what behavior cloning lets us skip). We drive it a few hundred times# and just remember what it saw and did. No dataset file, no disk: two arrays.def collect_demos(n_episodes: int, seed: int): env = PushTEnv() observations, actions, expert_successes = [], [], 0 for i in range(n_episodes): obs = env.reset(seed + i) # seed+i per episode: each demo is a different, reproducible start expert = ScriptedExpert(noise=0.0, seed=seed + i) done, info = False, {} while not done: action = expert.action(env) observations.append(obs.copy()) # the obs we acted ON, never the terminal one actions.append(action.copy()) obs, _, done, info = env.step(action) expert_successes += bool(info["success"]) return np.asarray(observations, np.float32), np.asarray(actions, np.float32), expert_successes demo_obs, demo_actions, expert_successes = collect_demos(args.demos, args.seed)print(f"demos: {args.demos} expert episodes ({expert_successes} solved) -> {len(demo_obs)} (obs, action) pairs")A tiny network. Three linear layers (a linear layer is just a matrix multiply plus an added offset). Its whole job is to map "the world looks like this" to "so do that". Boring on purpose.
# Rescale every input and output dimension to [-1, 1] using its min/max over the# demos — meters, sin/cos and velocities all live on different scales, and a# network learns far faster when they don't. The stats live INSIDE the model so# the policy is self-contained. (ch1.1 shows how a lie in these stats sinks a# policy that every loss curve swears is healthy.)def minmax(x): # -> (min, range) with zero-width dims (the constant target pose) held safe lo = x.min(0) span = np.where(x.max(0) - lo < 1e-4, np.float32(1.0), x.max(0) - lo) return torch.from_numpy(lo), torch.from_numpy(span.astype(np.float32)) class BCPolicy(nn.Module): """obs float32[10] -> action float32[2]. A 3-layer MLP: no planner, no search, no idea what a "T" is. It maps "when the world looks like this" to "do that", and that is enough to push a block home most of the time.""" 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, stat) # travels with the weights, never trained def forward(self, obs): normalized = (2.0 * (obs - self.obs_min) / self.obs_range - 1.0).clamp(-1.0, 1.0) action = self.net(normalized) return (action + 1.0) / 2.0 * self.act_range + self.act_min obs_min, obs_range = minmax(demo_obs)act_min, act_range = minmax(demo_actions)policy = BCPolicy(args.hidden_dim, obs_min, obs_range, act_min, act_range).to(device)Training. Show it a demonstrated state, ask what it would do, penalize the gap to what the expert actually did, repeat. That's the entire method. The loop is short enough to read in one sitting.
# The entire training method: predict the action the expert took, penalize the# squared error, repeat. No DataLoader, no early stopping — the demos are two# tensors in memory and a "batch" is a slice of a shuffled index permutation.train_obs = torch.from_numpy(demo_obs).to(device)train_actions = torch.from_numpy(demo_actions).to(device)optimizer = torch.optim.Adam(policy.parameters(), lr=1e-3)scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) # decay to 0 so the last epochs settleshuffle = torch.Generator().manual_seed(args.seed)for epoch in range(args.epochs): epoch_loss, num_batches = 0.0, 0 for batch in torch.randperm(len(train_obs), generator=shuffle).split(256): loss = nn.functional.mse_loss(policy(train_obs[batch]), train_actions[batch]) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss, num_batches = epoch_loss + loss.item(), num_batches + 1 scheduler.step() if args.rerun: rr.set_time("epoch", sequence=epoch) rr.log("policy/loss", rr.Scalars([epoch_loss / num_batches])) if epoch % 100 == 0 or epoch == args.epochs - 1: print(f"epoch {epoch:3d} loss {epoch_loss / num_batches:.5f}")The honest test. Loss on the demonstrations tells you the network memorized them; it does not tell you the policy can push a block. So we turn it loose on 25 starting positions it never trained on and count the ones where the block actually reaches the target. And we roll out a random policy on the same 25 starts as a floor. If "it works" is going to mean anything, it has to beat flailing.
# The only test that counts: turn the policy loose on starts it never trained# on and see whether the block reaches the target. Eval seeds are offset by# 10_000 so they can't collide with a demo start. We also roll out a RANDOM# policy on the same starts — the floor the trained policy has to clear for# "it works" to mean anything. And we save the first success for the browser.def rollout(act_fn, seed, record=False): env = PushTEnv() obs, done, info = env.reset(seed), False, {} frames = [] while not done: action = act_fn(obs) if record: px, py = env.pusher_pos tx, ty, tyaw = env.tee_pose pos_err, _ = env._errors() frames.append([round(float(v), 5) for v in (px, py, tx, ty, tyaw, pos_err)]) obs, _, done, info = env.step(action) return bool(info["success"]), frames policy.eval() def policy_action(obs): with torch.no_grad(): return policy(torch.from_numpy(obs).to(device).unsqueeze(0))[0].cpu().numpy() def random_action(_obs): return rng.uniform(-1.0, 1.0, size=PushTEnv.ACT_DIM).astype(np.float32) successes, saved_rollout = 0, Nonefor episode in range(args.eval_episodes): solved, frames = rollout(policy_action, 10_000 + args.seed + episode, record=True) successes += solved if solved and saved_rollout is None: # keep the first honest success for replay saved_rollout = {"seed": 10_000 + args.seed + episode, "frames": frames} if args.rerun: rr.set_time("eval_episode", sequence=episode) rr.log("eval/success", rr.Scalars([float(solved)]))random_successes = sum(rollout(random_action, 10_000 + args.seed + e)[0] for e in range(args.eval_episodes)) success_rate = successes / args.eval_episodesrandom_rate = random_successes / args.eval_episodesprint(f"\n trained policy: {successes}/{args.eval_episodes} solved ({success_rate:.0%})")print(f" random policy: {random_successes}/{args.eval_episodes} solved ({random_rate:.0%}) <- the floor you cleared")if success_rate > random_rate: # the real run: a visible, honest win over the random floor print("\nYou trained a robot policy and it works. Now go build every piece of it, starting in ch0.1.")else: # only the tiny --smoke budget lands here; it exists to check the plumbing, not to win print("\n(smoke budget — too small to solve the task; the full run is where it works.)") metrics = { "seed": args.seed, "demos": args.demos, "epochs": args.epochs, "eval_episodes": args.eval_episodes, "expert_successes": expert_successes, "successes": successes, "success_rate": round(success_rate, 6), "random_successes": random_successes, "random_rate": round(random_rate, 6), "smoke": bool(args.smoke),}(args.out / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n")# The rollout the site replays: one real, held-out success (never cherry-picked# beyond "the first one"). None only if the tiny --smoke budget solved nothing.rollout_path = args.out / "rollout.json"rollout_path.write_text(json.dumps({ "success_rate": round(success_rate, 6), "random_rate": round(random_rate, 6), "target_pose": [round(float(v), 5) for v in PushTEnv.TARGET_POSE], "rollout": saved_rollout,}, indent=2) + "\n")print(f"\nmetrics: {args.out / 'metrics.json'} rollout: {rollout_path}")if args.rerun: print(f"recording: {args.out / 'quickstart.rrd'} — open it with: rerun {args.out / 'quickstart.rrd'}")At --seed 0 the trained policy solves 12 of 25 held-out starts (48%). The
random floor solves 0 of 25. That gap, from nothing to roughly half, is
the whole point of this chapter, and you produced it in a few minutes on a
laptop with no GPU. (It's seed-dependent: across the first few seeds the policy
lands between about a quarter and a half of the starts. Never zero, always well
clear of the floor.)
What we skipped, and where it comes back
This was the shortest honest path to a working policy, which means it cut corners you should know about, and every one of them is a real chapter ahead, not a hand-wave:
- The demonstrations came from a script. Real imitation learning starts with your demonstrations. In ch0.4 you drive the pusher yourself and record them, and you learn what a dataset physically is.
- The data lived in memory. We never wrote a dataset to disk. The actual format, the one real robots record into, is ch0.4 and ch0.5.
- We trained with plain squared error and didn't ask what breaks. That's behavior cloning, and it has failure modes that no loss curve will confess to. ch1.1 is this exact method, built properly, with those failures induced and measured, including why 48% and not 100%.
- The simulator was a black box. What
env.stepactually does (the physics loop you just trusted a few hundred thousand times) is ch0.1, the very next thing.
So the rest of Phase 0 is not new material piled on top of this. It is this, slowed down: now let's understand what you just did.
Exercise
One, in exercises/, and it's worth doing before you move on. All 300 of those
demonstrations succeeded at the task. So did the win really need 300, or would
a handful of perfect demos have done just as well? Commit to an answer, then let
the code settle it.
What's next
You have a policy that works and no idea why. Chapter 0.1 opens the simulator
you just called env.step on and shows you the loop underneath: the first of
the pieces that, put back together, are the thing now pushing a block across
your screen.