The problem
Chapter 1.2 left us at a wall we could name but not climb. Behavior cloning fits the average action, and at a state where two futures are equally good, the average of "go left" and "go right" is "drive straight into the middle." We curated the data until the wall was as low as data could make it, and the policy still topped out well short of the expert. The diagnosis at the end of 1.2 was blunt: it is the model class. A function that commits to one action per state, re-evaluated every 0.1 seconds, cannot represent a temporally extended plan.
The bimanual hand-off makes this concrete. Consider the instant the right arm arrives at the middle with the cube. The correct behavior is a little sequence: hold, wait for the left gripper to close, then open and retreat. A single-step policy sees only the current frame and has to re-derive that whole intention from scratch every step — and near the hand-off, tiny errors in "am I holding or releasing" compound into dropping the cube. Nothing about predicting one action at a time lets the policy say "I am three steps into a five-step release." That memory has to live somewhere, and in ACT it lives in the chunk.
Build
act.py is one file, about 380 lines, in six regions: setup, data, model,
train, eval, report. It generates scripted-expert demos of the cube transfer,
reshapes them into action chunks, builds a small transformer from scratch,
trains it with plain L1, and rolls it out with temporal ensembling. There is a
real transformer in here and nothing is imported to hide it — no transformers,
no timm, no attention you cannot read.
Setup
import argparseimport jsonimport shutilimport sysfrom pathlib import Path import numpy as npimport torchimport torch.nn as nnimport torch.nn.functional as F # Chapter artifacts run as loose scripts from the repo root; put the root on# sys.path so `curriculum.common` resolves (same pattern as ch1.1's bc.py).sys.path.insert(0, str(Path(__file__).resolve().parents[3])) from curriculum.common.device import banner, detect_device # noqa: E402from curriculum.common.envs.aloha_cube import AlohaCubeEnv, gen_demos # noqa: E402from curriculum.common.seeding import set_seed # noqa: E402 HEADS = 4 # attention heads (fixed; model_dim is the scale knob)ENC_LAYERS = 2 # encoder self-attention blocks over the 4 entity tokensDEC_LAYERS = 2 # decoder blocks (self-attn over queries + cross-attn to memory)OBS_DIM, ACT_DIM = AlohaCubeEnv.OBS_DIM, AlohaCubeEnv.ACT_DIM parser = argparse.ArgumentParser(description=__doc__)parser.add_argument("--data", type=Path, default=None, help="LeRobot aloha_cube dataset; omitted => generate --num_demos scripted demos")parser.add_argument("--out", type=Path, default=Path("outputs/ch1.3-act"))parser.add_argument("--chunk_size", type=int, default=8, help="K: actions predicted per forward pass. aloha episodes are ~27 steps, so keep K well under that (measured: K=8 beats K=16 here, both crush K=1)") # smoke: 4parser.add_argument("--model_dim", type=int, default=128) # T4: 256 | smoke: 16parser.add_argument("--num_demos", type=int, default=50) # T4: 200 | smoke: 4parser.add_argument("--epochs", type=int, default=400) # cpu-laptop: minutes | smoke: 3parser.add_argument("--batch_size", type=int, default=64)parser.add_argument("--lr", type=float, default=1e-3, help="peak Adam lr; cosine-decays to 0 over --epochs")parser.add_argument("--ensemble_m", type=float, default=0.1, help="temporal-ensembling decay: larger = concentrate weight on the OLDEST overlapping prediction (commit to earlier plans); ~0 = uniform (real ACT uses ~0.01)")parser.add_argument("--eval_episodes", type=int, default=25) # T4: 50 | smoke: 2 — few episodes is noisy (ch1.6)parser.add_argument("--seed", type=int, default=0, help="seeds demo generation, the init, and the shuffle")parser.add_argument("--break", dest="break_mode", choices=("no_chunk", "no_ensemble", "open_loop"), default=None, help="Break It: a real ACT misconception with a measured signature (see the eval region)")parser.add_argument("--device", choices=("cpu", "cuda", "mps"), default=detect_device()) # cpu: deterministic (statistical repro on GPU/mps)parser.add_argument("--smoke", action="store_true", help="tiny self-contained CPU run for CI; two runs must produce byte-identical metrics.json")parser.add_argument("--rerun", dest="rerun", action="store_true", default=True)parser.add_argument("--no-rerun", dest="rerun", action="store_false", help="skip .rrd recording (CI smoke)")args = parser.parse_args() rng = set_seed(args.seed) # seeds python/numpy/torch; returns the numpy Generator demo-gen draws fromif args.smoke: # smoke pins everything the CI byte-compare depends on args.chunk_size, args.model_dim, args.num_demos = 4, 16, 4 args.epochs, args.eval_episodes, args.device = 3, 2, "cpu"if args.break_mode == "no_chunk": args.chunk_size = 1 # the whole ACT idea, ablated: predict a single action, like 1.1's BCbanner("ch1.3-act", device=args.device)args.out.mkdir(parents=True, exist_ok=True)device = torch.device(args.device)if args.rerun: import rerun as rr rr.init("zero2robot/ch1.3-act", spawn=False) rr.save(str(args.out / "act.rrd"))The scale knobs are the story of the chapter as flags. --chunk_size (K) is how
far ahead the policy commits; --model_dim is the width of the transformer.
Every default is free-tier first — the whole run fits on a CPU laptop in a few
minutes. A note the env forces on us: aloha_cube episodes are short, a median
of about 27 control steps, because the env is a deliberately tame planar model
with a weld-constraint grasp (see its README). So keep K well under the episode
length — a chunk longer than the task is mostly padding. Eight (about a third of
a median episode) is the default here, and it is not arbitrary: measured, K=8
beats K=16, which in turn crushes K=1 — the "tune the chunk down for short
episodes" rule, as a number. One hundred, which real ACT uses on its
several-hundred-step episodes, would be absurd on 27.
Data
# ACT trains on the SAME scripted-expert demos as everything else in Phase 1;# the only new step is reshaping per-frame actions into per-frame CHUNKS.if args.data is None: # Regenerate every run (never reuse a leftover dir): a cache from a different # --seed / --num_demos would silently train on the wrong data. gen_demos is # deterministic, so same args -> bit-identical demos whether built or rebuilt. args.data = args.out / "demos" if args.data.exists(): shutil.rmtree(args.data) gen_demos.main(["--episodes", str(args.num_demos), "--seed", str(args.seed), "--out", str(args.data), "--no-video"])if not (args.data / "meta" / "info.json").is_file(): sys.exit(f"no dataset at {args.data} — generate one first:\n" f" python curriculum/common/envs/aloha_cube/gen_demos.py " f"--episodes 50 --seed 0 --out {args.data} --no-video") from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: E402 (heavy import — after cheap failures) frames = LeRobotDataset("local/aloha_cube", root=args.data).hf_dataset.with_format("numpy")obs = np.stack(frames["observation.state"]).astype(np.float32) # (N, 10) — layout in aloha_cube_env.pyactions = np.stack(frames["action"]).astype(np.float32) # (N, 6) — already clipped to [-1, 1]episode_ids = np.asarray(frames["episode_index"]) # (N,) which demo each frame came from # Chunk targets: for frame i, the next K expert actions within ITS episode. Near# an episode's end there are fewer than K left, so pad by repeating the last# action and record a 0/1 mask so those padded steps carry no gradient. This is# what turns a per-step dataset into a per-chunk one (real ACT does the same).K = args.chunk_sizechunk_targets = np.zeros((len(obs), K, ACT_DIM), dtype=np.float32)chunk_mask = np.zeros((len(obs), K), dtype=np.float32)for e in np.unique(episode_ids): idx = np.nonzero(episode_ids == e)[0] # frame indices of this episode, in order ep_actions = actions[idx] for j, frame in enumerate(idx): valid = min(K, len(idx) - j) chunk_targets[frame, :valid] = ep_actions[j:j + valid] chunk_targets[frame, valid:] = ep_actions[-1] # pad (masked out below) chunk_mask[frame, :valid] = 1.0 # Normalization stats over obs only; actions are already in [-1, 1] by the env's# action contract, so the chunk head regresses them directly (no action denorm).obs_min = obs.min(0)obs_range = np.where(obs.max(0) - obs_min < 1e-4, np.float32(1.0), obs.max(0) - obs_min)obs_t = torch.from_numpy(obs).to(device)chunk_t = torch.from_numpy(chunk_targets).to(device)mask_t = torch.from_numpy(chunk_mask).to(device)print(f"dataset: {len(np.unique(episode_ids))} episodes / {len(obs)} frames, " f"chunk_size={K}, model_dim={args.model_dim}")The demos are the same scripted expert every Phase-1 chapter uses; the only new
step is the reshape. For each frame t, the target is no longer one action —
it is the next K actions the expert took, actions[t : t+K]. Near the end of an
episode there are fewer than K actions left, so we pad the tail by repeating the
last action and record a 0/1 mask that zeroes those padded steps out of the loss.
Get that mask wrong — mark the padding as real — and you quietly train the policy
to predict invented actions at the end of every episode (that is exercise 3). The
observations themselves are unchanged from chapter 1.1: ten numbers, normalized
into the model as buffers so the checkpoint carries its own stats.
Model
# A transformer block is attention + a feed-forward net, each wrapped in a# residual connection with pre-norm (LayerNorm before the sublayer — the stable# variant). We hand-roll it from nn.MultiheadAttention so every piece is visible;# nothing here is imported from a transformer library.class EncoderBlock(nn.Module): """Self-attention over the entity tokens, then a per-token FFN.""" def __init__(self, dim: int): super().__init__() self.attn = nn.MultiheadAttention(dim, HEADS, batch_first=True) self.norm1, self.norm2 = nn.LayerNorm(dim), nn.LayerNorm(dim) self.ff = nn.Sequential(nn.Linear(dim, 2 * dim), nn.GELU(), nn.Linear(2 * dim, dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: # (B, tokens, dim) h = self.norm1(x) x = x + self.attn(h, h, h, need_weights=False)[0] # tokens attend to each other return x + self.ff(self.norm2(x)) class DecoderBlock(nn.Module): """Query tokens self-attend, then cross-attend to the encoder memory, then FFN.""" def __init__(self, dim: int): super().__init__() self.self_attn = nn.MultiheadAttention(dim, HEADS, batch_first=True) self.cross_attn = nn.MultiheadAttention(dim, HEADS, batch_first=True) self.norm1, self.norm2, self.norm3 = nn.LayerNorm(dim), nn.LayerNorm(dim), nn.LayerNorm(dim) self.ff = nn.Sequential(nn.Linear(dim, 2 * dim), nn.GELU(), nn.Linear(2 * dim, dim)) def forward(self, q: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: h = self.norm1(q) q = q + self.self_attn(h, h, h, need_weights=False)[0] # chunk steps coordinate h = self.norm2(q) q = q + self.cross_attn(h, memory, memory, need_weights=False)[0] # read the scene return q + self.ff(self.norm3(q)) class ACTPolicy(nn.Module): """obs float32[B,10] -> action chunk float32[B,K,6]. The whole architecture. The obs is not a sequence, so we MAKE one: split the 10 numbers into four entity tokens (right arm, left arm, cube, target), each padded to width 3 and projected to model_dim. Self-attention over those four tokens is where the policy reasons about relationships — where is the cube relative to each gripper, which arm should act. K learned query tokens then cross-attend to that memory; each query becomes one action in the chunk. """ def __init__(self, dim: int, chunk_size: int, obs_min, obs_range): super().__init__() self.token_proj = nn.Linear(3, dim) # shared across the 4 tokens self.type_embed = nn.Parameter(torch.zeros(1, 4, dim)) # "which entity am I" self.query_embed = nn.Parameter(0.02 * torch.randn(1, chunk_size, dim)) # "which step am I" self.encoder = nn.ModuleList(EncoderBlock(dim) for _ in range(ENC_LAYERS)) self.decoder = nn.ModuleList(DecoderBlock(dim) for _ in range(DEC_LAYERS)) self.head = nn.Linear(dim, ACT_DIM) # each query -> one 6-D action for name, stat in [("obs_min", obs_min), ("obs_range", obs_range)]: self.register_buffer(name, torch.from_numpy(stat)) # saved with the weights, never trained def forward(self, observation: torch.Tensor) -> torch.Tensor: obs_n = (2.0 * (observation - self.obs_min) / self.obs_range - 1.0).clamp(-1.0, 1.0) tokens = torch.stack([ # (B, 4, 3) entity tokens obs_n[:, 0:3], # right arm: x, y, grip obs_n[:, 3:6], # left arm: x, y, grip F.pad(obs_n[:, 6:8], (0, 1)), # cube: x, y, (pad) F.pad(obs_n[:, 8:10], (0, 1)), # target: x, y, (pad) ], dim=1) memory = self.token_proj(tokens) + self.type_embed for block in self.encoder: memory = block(memory) queries = self.query_embed.expand(observation.shape[0], -1, -1) for block in self.decoder: queries = block(queries, memory) return self.head(queries) # (B, K, 6) policy = ACTPolicy(args.model_dim, K, obs_min, obs_range).to(device)This is the part worth slowing down for. A transformer attends over a set of tokens, but our observation is a flat vector of ten numbers — so we make a set out of it. We split the ten numbers into four entity tokens: the right arm (x, y, grip), the left arm, the cube, and the target. Self-attention over those four tokens is exactly where the policy gets to reason about relationships — how far is the cube from each gripper, which arm is in position — instead of staring at ten anonymous floats. That is the encoder.
The decoder is where the chunk comes from. We hold K learned query tokens,
one per step of the chunk. Each query first attends to the other queries (so the
steps of the plan can coordinate — "if step 3 releases, step 4 retreats") and
then cross-attends to the encoder's memory of the scene. A linear head turns
each finished query into one six-dimensional action. Eight queries in, eight
actions out, in a single forward pass. The blocks are hand-rolled from
nn.MultiheadAttention with pre-norm residuals — the same structure as the DETR
decoder real ACT borrows, shrunk until you can read all of it.
Train
# Plain loop, one policy, no DataLoader. The loss is L1 (ACT's choice — sharper# than MSE on multimodal action data) between the predicted chunk and the K# expert actions, averaged over the VALID (unpadded) steps only.def chunk_l1(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: per_step = (pred - target).abs().mean(-1) # (B, K) mean over the 6 action dims return (per_step * mask).sum() / mask.sum() # ignore padded steps optimizer = torch.optim.Adam(policy.parameters(), lr=args.lr)# Decay the lr to 0 over the run so the last epochs settle instead of bouncing.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)shuffle = torch.Generator().manual_seed(args.seed) # torch-side RNG: same seed -> same batch ordertrain_loss, global_step = float("nan"), 0for epoch in range(args.epochs): epoch_loss, num_batches = 0.0, 0 for batch in torch.randperm(len(obs_t), generator=shuffle).split(args.batch_size): loss = chunk_l1(policy(obs_t[batch]), chunk_t[batch], mask_t[batch]) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss, num_batches = epoch_loss + loss.item(), num_batches + 1 if args.rerun: rr.set_time("step", sequence=global_step) rr.log("policy/loss/train", rr.Scalars([loss.item()])) global_step += 1 scheduler.step() train_loss = epoch_loss / num_batches if epoch % 50 == 0 or epoch == args.epochs - 1: print(f"epoch {epoch:4d} chunk_l1 {train_loss:.5f}")The loop is chapter 1.1's, almost unchanged: Adam, cosine-decayed lr, a shuffled permutation for batches. The one difference is the loss. We use L1 on the chunk (real ACT's choice — it is sharper than MSE on multimodal action data), averaged over the valid, unpadded steps. That is the entire training story: map each observation to the expert's next eight actions, and measure the miss in absolute value.
Eval — temporal ensembling
# Loss measured imitation on dataset states; rollouts measure the task. At each# step the policy predicts a fresh chunk, but every step is ALSO covered by# chunks predicted at earlier steps. Temporal ensembling averages all of those# overlapping predictions for the current step with exponential weights (older# predictions weighted exp(-m * age)), so the executed trajectory is smooth and# committed rather than jerking at chunk boundaries. --break ablates this:# no_chunk K=1 (set in setup): single-step BC through a transformer — the# chunking benefit is gone; measure how far success falls.# no_ensemble execute only the freshly predicted chunk's first action, no# averaging across overlapping chunks.# open_loop commit to the whole chunk (run all K, then re-query): the naive# "why bother ensembling" version, jerky at every chunk seam.MAX_T = AlohaCubeEnv.MAX_STEPS @torch.no_grad()def predict_chunk(net: ACTPolicy, observation: np.ndarray) -> np.ndarray: obs_batch = torch.from_numpy(observation).to(device).unsqueeze(0) # (10,) -> (1, 10) return net(obs_batch)[0].cpu().numpy() # (K, 6) def rollout(net: ACTPolicy, seed: int, mode: str, tag: str, episode: int) -> tuple[bool, float]: env = AlohaCubeEnv() observation = env.reset(seed) all_time = np.zeros((MAX_T, MAX_T + K, ACT_DIM), dtype=np.float32) # [query time, target time] populated = np.zeros((MAX_T, MAX_T + K), dtype=bool) committed: list[np.ndarray] = [] done, episode_return, info, t = False, 0.0, {}, 0 while not done and t < MAX_T: if mode == "open_loop": if not committed: # re-query only when the plan runs out committed = list(predict_chunk(net, observation)) action = committed.pop(0) else: chunk = predict_chunk(net, observation) all_time[t, t:t + K] = chunk populated[t, t:t + K] = True if mode == "no_ensemble": action = chunk[0] else: # temporal ensembling: average every chunk that predicted step t votes = all_time[:t + 1, t][populated[:t + 1, t]] # oldest first weights = np.exp(-args.ensemble_m * np.arange(len(votes))) action = (votes * (weights / weights.sum())[:, None]).sum(0) observation, reward, done, info = env.step(action) episode_return += reward t += 1 if args.rerun: rr.set_time("sim_time", duration=episode * (MAX_T / AlohaCubeEnv.CONTROL_HZ) + env.data.time) rr.log(f"eval/{tag}/action", rr.Scalars(action.astype(np.float64))) rr.log(f"eval/{tag}/dist", rr.Scalars([info["dist"]])) return bool(info["success"]), episode_return def evaluate(net: ACTPolicy, mode: str, tag: str) -> tuple[float, float]: outcomes = [rollout(net, 20_000 + args.seed + ep, mode, tag, ep) # held-out seeds, never in demos for ep in range(args.eval_episodes)] success_rate = float(np.mean([s for s, _ in outcomes])) mean_return = float(np.mean([r for _, r in outcomes])) if args.rerun: rr.log(f"eval/{tag}/success_rate", rr.Scalars([success_rate])) print(f"eval[{tag:11s}]: success {success_rate:.2f} mean_return {mean_return:.3f}") return success_rate, mean_return eval_mode = args.break_mode if args.break_mode in ("no_ensemble", "open_loop") else "ensemble"baseline_success, baseline_return = evaluate(ACTPolicy(args.model_dim, K, obs_min, obs_range).to(device), "ensemble", "untrained") # random-init referencesuccess_rate, mean_return = evaluate(policy, eval_mode, "trained")Here is the subtle idea, and it only exists because we chunk. At execution step
t, the policy predicts a fresh chunk covering steps t … t+K-1. But step t
was also predicted by the chunk from step t-1, and t-2, all the way back K
steps — each of those older chunks reached forward and made a prediction for
right now. Temporal ensembling averages all of those overlapping predictions
for the current step, weighting them exponentially so the policy blends its
plans into one smooth, committed trajectory instead of lurching each time a fresh
chunk disagrees with the old one at the seam. It is a running vote across time,
and it costs nothing but a buffer.
The --break flag ablates each piece so you can feel what it was worth; the
"Break it" section below reads the numbers.
Run it
python curriculum/phase1_imitation/ch1.3_act/act.py --seed 0 --device cpu
| cpu-laptop | expected wall-clock on cpu-laptop: ~5.44 min (measured) | measured |
| mps | wall-clock on mps: not yet measured | pending |
| t4 | expected wall-clock on t4: ~3.42 min (measured) | measured |
| 4090 | wall-clock on 4090: not yet measured | pending |
The result at the default config, seed 0 on CPU, over 25 held-out eval episodes:
| held-out success | mean return | |
|---|---|---|
| untrained (random init) | 0% | −320 |
| trained ACT (chunk + ensembling) | 88% | −46 |
The network never touched the environment during training — it only ever saw the expert's chunks — and it comes out driving the cube through the hand-off on nearly 9 of 10 held-out starts. That gap is the chunking and the transformer doing their job. Open the recording and scrub the two arms through the hand-off band:
rerun outputs/ch1.3-act/act.rrd
It is not the expert's 100%, and it should not be: this is a tiny transformer on fifty demonstrations trained for a few minutes so it fits on a laptop. Scale the demos, the width, and the epochs and the number climbs — but at this budget the chunked policy already clears a bar single-step behavior cloning never reached on this task.
What we cut
This is a real transformer trained the real way, but it is not the full ACT, and the missing pieces matter enough to name:
- No CVAE. Real ACT is a conditional variational autoencoder: a second
encoder reads the expert's action sequence and compresses it into a latent
zthat the decoder conditions on, trained with a KL term. That latent is how real ACT models multimodality — a demonstrator who sometimes goes left and sometimes right. We dropped it entirely. Our policy is deterministic: obs in, one chunk out. It works here because the scripted expert is essentially unimodal (one clean way to do the hand-off), so there is little multimodality for a latent to capture. On messy human demos, the missing CVAE is exactly what you would add back. - No images. Real ACT sees the scene through cameras and a ResNet backbone.
We train on the ten-number state vector, so there is no vision at all — the
entity tokens are our perception. The env can render a top-down image
(
--videodemos), and wiring an image encoder onto the front of this same transformer is the natural next step.
Neither cut is an approximation that quietly degrades a number; each is a whole capability left for later, on purpose, so the chunk-and-ensemble core is legible. The "read the real thing" segment for this chapter walks the original ACT repo so you can see precisely what these two paragraphs left out.
Break it
Three ablations, each a real ACT misconception, all measured at the default config (seed 0, CPU, 25 eval episodes). The first has a large, robust signature; the other two teach something subtler and more honest about what temporal ensembling actually buys.
--break no_chunk — "why not just predict one action?" This forces K=1: the
transformer now predicts a single action, which is chapter 1.1's behavior cloning
wearing an attention costume. Held-out success drops from 0.88 to 0.6 — and
the gap holds across seeds and eval sizes. The chunk was not decoration; it was
the mechanism. A policy that re-decides everything every step cannot hold the
hold-wait-release intention through the hand-off, and the cube gets dropped. This
is the chapter's thesis stated as an ablation, and it is solid.
--break no_ensemble — "the chunk is enough, skip the averaging." Same
trained weights, but at eval we execute only the first action of each fresh chunk
and throw the overlap away. Now be careful, because the honest result is more
interesting than a clean win. On seed 0 the policy collapses to 0.0 — total
failure — while on seed 1 it holds at 0.96, as high as it ever gets there
(seed 1 with ensembling is 0.88, so dropping it cost nothing — if anything the
sign flipped, well inside the noise of 25 episodes). Temporal ensembling did not
reliably raise success; on one seed it was the difference between everything and
nothing, on another it did nothing at all. The mechanism is
the env's honest simplification biting back: the grasp is a threshold weld
(close past CLOSE_FRAC while within reach of the cube and it latches). The
gripper command sits right at that threshold during the pick, and the temporal
average of several confident chunks is sometimes exactly what tips it over to
latch. Strip the averaging and, on an unlucky seed, the grasp never catches. What
ensembling reliably gives you is not success but smoothness — scrub
eval/*/action in the two recordings and the ensembled trace is smooth where the
no_ensemble one chatters. Its designed payoff, robustness to real-robot
observation noise, is a thing this clean deterministic sim cannot show you at all.
--break open_loop — "commit to the whole chunk, then look again." Run all
eight predicted actions, then re-query. Here it lands at 0.88, the same as
ensembling — on this task the jerk at chunk seams is visible in the traces but not
fatal to success. It is the cleanest illustration that "chunk without ensembling"
is a real design point, not obviously wrong, and that the smoothness argument for
ensembling is an argument about trajectories, not always about task success.
The transferable lesson, and it is a sharper one than "ensembling is good": chunking and ensembling are two separate ideas bundled under one name. Chunking buys you a plan, and here it robustly buys success. Ensembling buys you a smooth plan; whether smoothness also buys success depends on the task, the seed, and how marginal your contacts are — and measuring that honestly, instead of assuming it, is the whole game (chapter 1.6 is about exactly this fragility).
Read the real thing
You have now built every idea real ACT is made of except the two we cut, and the
original code is public — tonyzhaozh/act, pinned here at commit 742c753. Read it
next to act.py and the shape of what we simplified becomes exact.
Temporal ensembling. Our eval region keeps an all_time buffer, marks which
past chunks reached step t with a populated mask, and averages them with
weights = np.exp(-args.ensemble_m * np.arange(len(votes))). The real loop is in
imitate_episodes.py, inside eval_bc (around lines 219–260):
all_time_actions[[t], t:t+num_queries] = all_actions, then actions_for_curr_step = all_time_actions[:, t], a "populated" check via torch.all(... != 0), and k = 0.01; exp_weights = np.exp(-k * np.arange(...)). This is the one place where our
code is nearly line-for-line the original. The only real differences: theirs runs on
GPU against camera images, gates the whole thing behind a temporal_agg flag that
defaults off (you opt in, with query_frequency set to 1), and uses k=0.01
where our default is 0.1. If you read one file to confirm we did not fake the eval,
read this one.
The CVAE we cut. Our ACTPolicy is deterministic: obs in, one chunk out, no
latent anywhere. The real policy's loss, in policy.py (ACTPolicy.__call__,
lines 18–35), is l1 + kl_weight * kl — a second term we never compute. The
machinery lives in detr/models/detr_vae.py: a style encoder (encoder_action_proj,
cls_embed, ~lines 69–104) reads the expert's whole action sequence, latent_proj
splits it into mu/logvar, and reparametrize (lines 17–20) samples a 32-D latent
z the decoder conditions on. At eval, z is zeroed to the prior mean (lines
112–113). That entire path is how real ACT represents a demonstrator who sometimes
goes left and sometimes right — the multimodality our unimodal scripted expert
doesn't have, which is why we dropped it whole rather than approximated it.
The DETR transformer. Our EncoderBlock / DecoderBlock (self-attention, then
cross-attention to memory, pre-norm) mirror detr/models/transformer.py —
TransformerEncoderLayer and TransformerDecoderLayer, with the same self_attn +
multihead_attn cross-attention and the forward_pre / forward_post norm variants
we chose pre-norm from. Real ACT runs 4 encoder / 7 decoder layers at 8 heads (the
ACT branch of imitate_episodes.py: enc_layers = 4, dec_layers = 7, nheads = 8) where we run 2 / 2 at 4. And it feeds image tokens: detr/models/backbone.py
is a ResNet-18 whose feature map becomes the token set, with sinusoidal
detr/models/position_encoding.py — the whole vision front-end our four entity
tokens stand in for.
None of these three make our version wrong; they make it a legible minimum, with the
production hardening peeled off so the chunk-and-ensemble core is the only thing on
the page. Read them in this order: imitate_episodes.py eval_bc first (you
already know it), then policy.py (short — find the l1 + kl loss), then
detr/models/detr_vae.py's forward (the CVAE encoder, the piece we never built),
and last detr/models/transformer.py (the blocks you hand-rolled, at full size).
Exercises
Four, in exercises/. Two ask you to commit to a prediction before the run is
allowed to answer — the chunked policy against the untrained baseline, and the
chunked policy against its own single-step (K=1) ablation. One is a bug-hunt where
the padding mask marks invented actions as real and every metric still prints
clean. One has you implement the temporal-ensembling weighting from its
definition, since it is the one line the whole eval turns on.
What's next
You now have a policy that predicts the future in chunks and executes it smoothly. But it still learns by imitation — its ceiling is the demonstrator, and it has never once been told whether an action was good, only whether it matched the expert. It cannot discover a hand-off the expert never showed it, and it cannot recover from a state no demonstration ever visited. The next phase stops handing the policy answers and starts handing it a reward: the policy tries, fails, and improves from its own experience. Everything you have built — the chunk, the transformer, the honest env — carries forward; what changes is where the learning signal comes from.