BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning

In plain words

Teaching a robot a new job usually means collecting demonstrations of that exact job. BC-Z asks how far a robot gets if it has learned from many different jobs instead. The authors recorded a robot arm doing 100 tasks, with a person steering it through a VR controller. While the learned policy ran, the person could also step in and correct it, and those corrections became more training data.

The robot is told which task to do in one of two ways: a short sentence, or a video of a person doing the task. The sentence is turned into a vector by a pretrained sentence encoder, and that vector conditions the network that maps camera images to motions. Because the policy learns to follow the command, it can attempt tasks it never saw during training, with no demonstrations of those tasks. Video commands work the same way in principle, but in these experiments they generalised much worse than sentences.

The point is that a plain imitation-learning setup, scaled to many tasks, can pick up new tasks from a language description. The numbers come from one real robot setup, so they show what is possible there, not a guarantee for other robots.

TL;DR

• BC-Z is a behaviour-cloning policy for a real arm: camera image plus task embedding in, a short trajectory of end-effector motions out §3 §5.2.
• Training data: 25,877 robot demonstrations over 100 tasks, including HG-DAgger interventions, plus 18,726 human videos §1 §4.
• Author-reported: language-conditioned success averages 32% over the held-out tasks; video-conditioned success is 4% tab. 2 tab. 3.
• The released TensorFlow code covers only the language-conditioned model, and its training settings differ from the paper in several places (see Paper vs code).

Problem

The goal is a visuomotor policy that performs new tasks, given a command, without demonstrations of those tasks §1 §4. The authors argue that zero-shot generalisation to new tasks remains hard for vision-based manipulation, while prior imitation work mostly generalised to new objects or object placements §1. Their working hypothesis is that scaling and broadening the data, plus flexible task conditioning, is what makes generalisation possible §1 §3. The paper also uses "few-shot" for video conditioning, where the test-time embedding is an average of a few human videos §6.2.

Method

Setup. The policy is a map from an RGB image and a task command to an action. The action is the 6-DoF end-effector pose plus one gripper dimension §3. The policy is split into an encoder that turns the command into a task embedding z, and a control layer that maps the image and z to the action §3.

Data collection. A VR teleoperation rig drives the robot at 10 Hz §4. Collection starts with expert-only demonstrations. Then a learned policy runs while the operator supervises. The operator takes over with an override control to fix mistakes. This is HG-DAgger: interventions are added to the data and the policy is retrained §4. The paper reports 16 deployment iterations in total §4. Each iteration deploys the latest policy trained on all data so far §4. Tasks are sampled uniformly at random per episode, because letting operators pick tasks biased the data app. B.

Intervention rate as a proxy. The number of interventions per episode correlates negatively with success, so it can be tracked during collection as a cheap progress signal §4 fig. 5. This is a correlation the paper observed during HG-DAgger runs; it is not a measure of held-out generalisation.

Language encoder. A pretrained multilingual universal sentence encoder maps each command to a 512-dim vector §5.1. The released model takes this vector as an input feature model.py.

Video encoder. A human video is sampled to 20 frames, arranged in a 4×5 grid, and passed through a 2D ResNet-18 app. E. After mean pooling, the features go through a 32-unit ReLU layer and a 512-unit linear layer, then are L2-normalised app. E. It is trained end to end with the behaviour-cloning loss through the policy, plus a language-regression term: a cosine distance between the video embedding and the language embedding of the same task §5.1 eq. 1. Robot demonstrations are also embedded with this encoder and pushed toward the language vector §5.1 app. E alg. 1.

\min \sum_{i} \sum_{\substack{(s,a)\sim\mathcal{D}_e^i \\ w_h\sim\mathcal{D}_h^i\cup\mathcal{D}_e^i}} \big[ -\log\pi(a\mid s,z^i) + D_{\cos}(z_h^i, z_\ell^i) \big] eq. 1

Here z_h = q(w_h) is the video embedding, z_ℓ = q(w_ℓ) the language embedding, and D_cos the cosine distance §5.1 app. E. The authors say the language term is what organises the video embedding by task rather than by scene §5.1 app. E app. F.

Policy network. A ResNet-18 torso reads the monocular image; its pooled output feeds one head per action part §5.3. Each head is an MLP with two hidden layers of 256 units and ReLU, predicting delta XYZ, delta axis-angle, and the gripper value §5.3. The 512-dim task embedding conditions the torso through FiLM layers: a linear projection gives channel-wise scales and shifts for each of the four ResNet blocks §5.3. fig. 3

Action targets. Actions are not the raw demonstrated commands. Each action is the difference from the current state to a target pose N steps ahead, where N is chosen per state by an adaptive rule §5.2 app. C. The rule begins at N = 1 and keeps raising N while the gripper and arm barely change; it stops once the gripper has moved by more than 0.01 or the norm of the arm joint change passes 0.05 app. C. Labels in which neither the end effector nor the gripper moves are dropped app. C.

Open-loop auxiliary outputs. The policy predicts 10 future actions but executes only the first one, closed loop §5.2. The other nine act as an auxiliary training signal §5.2.

Training step. The sketch below combines alg. 2, app. D and §5.2. It is written from the paper, not from the released code.

# Language task vector: 512-dim sentence embedding, L2-normalised  [alg. 2]
def task_vector(command):
    return l2_normalize(sentence_encoder(command))     # (512,)

# One training step on a batch of images, task vectors and targets
def training_step(images, task_vectors, targets):
    # Gaussian noise on the task embedding, training only  [app. D]
    conditions = task_vectors + normal(0.0, 0.1)       # (B, 512)
    # ResNet-18 torso, FiLM-conditioned in all four blocks  [§5.3]
    features = film_resnet18(images, conditions)       # (B, 512)
    # One head per action part: two hidden layers of 256, ReLU  [§5.3]
    xyz = head_xyz(features)                           # (B, 10, 3)
    axis_angle = head_axis_angle(features)             # (B, 10, 3)
    gripper = head_gripper(features)                   # (B, 10, 1)
    # Targets are state differences to poses N steps ahead  [§5.2] [app. C]
    # Huber (delta 1.0) on xyz and axis-angle, weights 100 and 10  [app. D]
    loss = 100.0 * huber(xyz, targets.xyz)
    loss += 10.0 * huber(axis_angle, targets.axis_angle)
    # Log loss on the normalised gripper angle, weight 0.5  [app. D]
    loss += 0.5 * log_loss(gripper, targets.gripper)
    return loss

Video-encoder training. Each batch samples tasks with replacement. For each task it draws one human video and one robot demonstration, so the batch is half human and half robot app. E alg. 1. The combined loss is the behaviour-cloning term plus two cosine terms, one for the human video and one for the robot video, both against the language vector alg. 1.

Evaluation with video. At test time the video task vector is the average of several new human videos of the task; the appendix says 10 app. D. The appendix also reports that averaging several video embeddings before conditioning worked better than not averaging app. D.

Hyperparameters

Results

Author-reported. Setup: real robot, 100 training tasks, ResNet-18 FiLM policy, 29 held-out tasks in the text (28 rows in the table). Baseline per row is noted in the table.

Single-task validation §6.1 tab. 1

Held-out tasks §6.2 tab. 2. Success on each task is in %, with 1 standard deviation in parentheses. Columns: language with one distractor object, language with 4–5 distractors, video with 4–5 distractors.

Other held-out results:

• Among the 24 held-out tasks with non-zero language success, the mean is 44% §6.2 [abstract]. Recomputing from the one-distractor column of tab. 2 gives 1070 / 24 ≈ 44.6%; this card reports the paper's 44%.
• Single-task BC trained on 300+ demos of each of the first four held-out tasks scored 0% on all four §6.2.
• Qualitatively, the main failure is "last-centimetre" errors, such as not closing or not releasing the gripper §6.2.

Encoder versus policy bottleneck §6.2 tab. 3

The authors read the near-equal one-hot and language numbers as evidence that the language space is adequate, so held-out failures sit in the policy rather than in the embedding §6.2. The drop for video suggests inferring a task from video is harder §6.2.

Ablations §6.3 tab. 4. All on "place the bottle in the ceramic bowl", the task with the most demos (1,000).

Other reported numbers

• Video-encoder batch sampling (no BC loss): task-based sampling 84%, 28 human plus 28 robot videos 80%, uniform over all videos 74% app. E.
• Single-task BC in simulation, one bottle-in-bowl task: 97.2% (0.7) with 37 demonstrations and no scene randomisation; 56% (2.2) with 40 demonstrations when objects are randomised app. B.
• Pretrained video embedding and TecNets embedding, on 7 selected held-out tasks: average 14.4% (video-conditioned BC-Z), 9.3% (pretrained video embedding), 8.6% (TecNets) app. M.
• Intervention rate versus success across HG-DAgger checkpoints: a clear negative correlation §6.3 fig. 5.

Paper vs code

Code: the official repository is google-research/tensor2robot, folder research/bcz, at commit 034f49c435168151930c6a07abea48b0df58b74b, Apache-2.0. The paper's project page links this folder (checked with WebFetch), and the folder's README points back to the paper README.md. The released code is a TensorFlow policy model and gin configs. It does not include the video encoder, the data generator, or the adaptive-N label code.

• Multi-task learning rate and batch. Paper: 5e-3 and batch 4096 on a TPUv3 pod app. D. Released config: LEARNING_RATE = 2.5e-4 and TRAIN_BATCH_SIZE = 32 run_train_bc_gtcond_trajectory.gin run_train_bc_gtcond_trajectory.gin. The config is the baseline the README uses, not the paper's run.
• Training length. Paper: not stated. Released config: max_train_steps = 50000 run_train_bc_gtcond_trajectory.gin.
• Dataset mixing. Paper: no sampling ratio between the two object sets is stated §6.2. Released language config: weights = [0.1, 0.9] with the comment "10% 21-task data, 90% 83-task data" run_train_bc_langcond_trajectory.gin. The paper says 79 tasks, not 83, and the README dataset names say 79 README.md.
• Gripper target. Paper: log loss on a normalised gripper angle, range 0 to 1 app. D. The language config turns binarisation off and rescales labels to max(0, (x − 0.2) / 0.8) run_train_bc_langcond_trajectory.gin model.py. At inference the prediction maps back as 0.2 + v · 0.8 model.py. The paper does not mention the 0.2 floor. The base preprocessor defaults to binarisation model.py, and the gtcond config also sets it to True run_train_bc_gtcond_trajectory.gin.
• Residual targets. Paper: actions are state differences to future poses §5.2. Code: losses are computed on *_residual labels model.py; the absolute pose is rebuilt at inference by adding the current pose model.py. This matches the paper, but the label generation and the adaptive N rule are not in this repository.
• Gradient through future waypoints. Paper: not stated. Code: the head can stop the gradient from the future waypoints into the torso, and the default is True bcz_networks.py. The gtcond config sets it to False run_train_bc_gtcond_trajectory.gin run_train_bc_gtcond_trajectory.gin. Appendix L says stopping this gradient was inconclusive app. L.
• Image size. Paper: random crop and downsample, no sizes §5.2. Code: crop 450, resize 150 in the gtcond config run_train_bc_gtcond_trajectory.gin; the language config overrides IMAGE_SIZE = 200 run_train_bc_langcond_trajectory.gin.
• Task embedding normalisation. Paper: task vectors are L2-normalised alg. 2. Code: the 512-dim sentence_embedding is consumed as given model.py. Any normalisation would have to be in the data pipeline, which is not in this repository.
• Video conditioning. Paper: a video encoder trained with a language-regression loss §5.1 alg. 1. Code: ConditionMode has only ONEHOT_TASKID and LANGUAGE_EMBEDDING model.py. I found no video encoder in research/bcz or layers/bcz_networks.py.
• FiLM form. Paper: scales and shifts per channel §5.3. Code: (1 + gamma) * x + beta resnet.py. This is the same family; the paper does not give the exact form.
• Optimiser. Paper: Adam app. D. Code: the optimiser factory is in tensor2robot's utils, not in this repository.
• Matches checked: ResNet-18 run_train_bc_gtcond_trajectory.gin; FiLM on all four blocks run_train_bc_gtcond_trajectory.gin; two-layer 256 heads run_train_bc_gtcond_trajectory.gin; 10 waypoints run_train_bc_gtcond_trajectory.gin; loss weights 100, 10, 0.5 run_train_bc_gtcond_trajectory.gin; task noise 0.1 added in training only model.py run_train_bc_langcond_trajectory.gin; 512-dim language input model.py.
• Adaptive N, idle-label dropping, and the video encoder are not in this repository; I searched research/bcz and layers/bcz_networks.py for them and found nothing.

Pitfalls

• Held-out count. The text says 29 held-out tasks §4 §6.2. The table has 28 rows, its overall row matches a mean over 28, and the caption of tab. 3 says 28 tab. 2 tab. 3. Treat 28 as the count the table supports.
• Trials per task are not stated. The numbers in parentheses are one standard deviation tab. 2, but the number of trials per held-out task is not given. With success rates that swing by 15 points between columns, small trial counts make most individual cells noisy.
• Video conditioning is weak and not in the release. Held-out video success is 4% tab. 3, and the released code cannot train the video encoder. Do not build on the video path from this card without the authors' code.
• Language conditioning is the only route the release can train. Even that route's config overrides the image size to 200 px, while the shared base config says 150 px (see Paper vs code).
• Deterministic unimodal policy on a multimodal task. Bin emptying is multimodal, yet the deterministic policy works, which the authors attribute to variety and noise in the data breaking symmetry app. H. The mechanism is a hypothesis, not a result.
• Gripper and residual conventions. Several conventions decide whether the policy works: state-difference targets, adaptive N app. C, discarded idle labels app. C, gripper rescaling (Paper vs code), and Huber versus MSE app. L. Each is a silent place to diverge from the paper.
• Checkpoint selection. Validation error on future poses correlates with success, but models with similar error can differ widely in success, so checkpoint choice is hard app. L.
• Intervention rate is a proxy. It correlates with success during HG-DAgger §6.3, but the paper uses it only as a development signal, not as a generalisation metric.
• Demonstration speed. The appendix says humans demonstrate with video 5–7 times faster than with teleoperation app. B.
• Single-task baselines are not like-for-like. The 5% single-task result uses 1,000 demonstrations from a dataset spread over several robots and operators; the authors' own simulation shows a single-task policy can reach 97% when variability is removed app. B.
• Sim-to-real for door opening uses RetinaGAN to adapt simulated images, trained on a mixture of sim, real and adapted-sim data app. I. The door results in tab. 1 depend on this, and the RetinaGAN setup is not described in enough detail to reproduce.
• Dataset licence. The demonstrations are published on Kaggle §1 README.md. The dataset licence was not checked for this card.
• Zero-shot versus few-shot. The paper uses both terms, depending on the conditioning. Zero-shot is the language route; the video route averages several human videos at test time §6.2 app. D.

Lineage

• Builds on: HG-DAgger, a human-gated DAgger for interactive imitation learning (Kelly et al., 2019) §4; FiLM conditioning (Perez et al., 2018) §5.3; task-embedded control networks for few-shot imitation (James et al., 2018) app. E; the multilingual universal sentence encoder, arXiv:1907.04307 §5.1.
• Compared against: single-task behaviour cloning, one-hot task conditioning, a pretrained instructional-video embedding (Miech et al., 2020), and TecNets tab. 3 tab. 4 app. M.
• Superseded by: none known from this wiki's content.