Behavior Cloning with Regression from Absolute Zero

What You'll Master

The Flappy Bird Setup

You've been given an environment that's a physics-based version of the classic Flappy Bird game. A bird falls under gravity. Pipes scroll left across the screen. Each pipe has one or two gaps. The bird must navigate through the gaps without crashing.

The agent's job: at each timestep, output a target y-position in [0, 1]. A built-in PD controller converts that target into thrust, creating realistic momentum-based motion.

The observation

4 numbers, all normalised to [0, 1]:

• obs[0]: distance to the next pipe
• obs[1]: y-position of gap 1
• obs[2]: y-position of gap 2 (in easy mode, same as gap 1)
• obs[3]: bird's current y-position

Two difficulty modes

The "hard" mode is where the lesson lives. With double-gap pipes, the expert can choose either gap. Different demonstrations choose differently. We'll see in Chapter 07 why this breaks MSE regression.

Action chunking

Modern imitation-learning policies don't predict one action at a time. They predict a chunk of T future actions in a single forward pass:

In this homework: ACTION_CHUNK = 20, EXECUTE_STEPS = 10. The policy predicts 20 future targets, but only the first 10 are executed before re-querying. This is called receding horizon control. We'll explain why this helps in Chapter 04.

Evaluation

Average episode length over 50 evaluation episodes. Episodes cap at 1000 steps; surviving all 1000 is "success." Higher is better. Standard deviation tells you about consistency — high std means the policy sometimes succeeds and sometimes immediately crashes.

The Imitation Learning Paradigm

Imitation learning is fundamentally different from reinforcement learning. If you're coming to HW1 fresh, internalize this distinction before anything else.

The imitation learning setup:

1. An expert exists. We have access to their demonstrations: a dataset of (state, action) pairs.
2. We want to train a policy π_θ(s) that produces actions like the expert would.
3. We don't ever query the expert's reward function or interact with the environment during training. (Problem 3, DAgger, slightly relaxes this last constraint.)

That's it. No exploration. No reward function. No Bellman equation. Just supervised learning with state → action pairs. The complexity comes from the subtleties: which loss function, which architecture, and what to do when the expert is multimodal.

Why imitation learning matters

Three reasons it's the right starting point for a course on robot learning:

• It's the simplest baseline. Every advanced method (DAgger, GAIL, IQL, RLHF) builds on top of BC.
• It's how most real robot policies are bootstrapped. Tesla's autopilot, OpenAI's dactyl, Stanford ALOHA, all start from BC pretraining before any RL fine-tuning.
• It exposes the deep questions about distribution shift, multimodality, and inductive biases that the rest of the course will solve.

Behavior Cloning

Strip BC down to its essence and it's literally just supervised learning.

The dataset

Run the expert in the environment N times. Record every (state, action) pair you see. Result: a dataset

where each ai* is the action the expert took at state si.

The objective

Find policy parameters θ that minimize the discrepancy between the policy's predicted action and the expert's action, averaged over the dataset:

where L is some loss function. For Problem 1, L is mean squared error. (Problem 2 will use flow matching, a generative loss. Different loss, same overall structure.)

What this means in practice

Three steps:

1. Collect demos: run the expert, log states and actions.
2. Define the model: a neural network π_θ(s) → a.
3. Optimize: standard supervised learning loop. Sample minibatches, compute loss, backprop, step Adam, repeat.

If you've trained any classifier or regressor in PyTorch before, you've already done all the mechanics of BC. The only thing that's new is the data — states and actions instead of images and labels.

Why BC isn't trivially perfect

The hardness of BC isn't in the optimization — it's in:

1. Distribution shift (DAgger fixes this in Problem 3): at test time the policy makes small errors that compound, and pretty soon it's in a state the expert never visited. The policy has no idea what to do.
2. Multimodal experts (Flow matching fixes this in Problem 2): if the expert sometimes goes through gap 1 and sometimes through gap 2, MSE regression averages the two and produces a midpoint — which crashes into the wall.
3. Causal confusion: the policy might learn shortcuts that work on the training data but don't generalize.

Problem 1 of HW1 introduces BC. Problem 2 (flow matching) tackles multimodality. Problem 3 (DAgger) tackles distribution shift. Each problem is one BC failure mode and one fix.

Why Action Chunking

Most BC tutorials predict one action at a time:

Modern robot learning (e.g., Diffusion Policy, ALOHA, RT-1) predicts a chunk of future actions:

where T = 20 in this homework (the ACTION_CHUNK constant). At rollout time, the policy is queried, then the first EXECUTE_STEPS = 10 actions are executed open-loop (without re-querying), and only after the 10th step does the policy fire again.

This is called receding horizon control — the policy commits to a 20-step plan, executes 10 steps of it, then re-plans.

Why does this help?

Three big reasons:

1. Temporal consistency

If you predict one action at a time, the policy can flip-flop — this step it predicts "go up," the next step "go down." Action chunks force temporal coherence: the network has to commit to a continuous plan, not a shaky sequence of independent decisions.

2. Reduced compounding error

Each policy query is one source of stochasticity. Querying once every 10 steps means you accumulate fewer "rolls of the policy dice" per episode. If errors are roughly Gaussian-distributed per query, fewer queries = lower total error variance.

3. Multi-step reasoning

Predicting 20 steps ahead lets the policy plan around obstacles (e.g., commit to a gap, route through it). A single-step policy has to make this decision fresh at every step, which gets tangled in noise.

How chunking shapes the architecture

For BC with action chunking on this homework:

• Input: state vector of shape [B, 4].
• Output: chunk of action targets of shape [B, 20].

The policy is just an MLP that maps 4 → 256 → 256 → 20. Each output dimension corresponds to one future action. The action_dim parameter in BCPolicy defaults to 20 because that's the chunk size.

The MLP Architecture

Problem 1 specifies a 3-layer MLP:

Three things to understand here.

Why an MLP and not something fancier

The problem is small (4-D input, 20-D output, simple physics). Convolutional networks would be overkill (no spatial structure in the state). Transformers would be overkill (no sequential structure in the input). A plain MLP with two hidden layers of 256 is the natural choice.

For comparison, the FlowMatchingPolicy for Problem 2 is much fancier — a 1D U-Net with conditional residual blocks. That's because flow matching needs to learn a vector field, which is a more complex function class.

Why ReLU

ReLU(x) = max(0, x). Three reasons to use it:

• Cheap to compute (one comparison + one multiply).
• Doesn't saturate for positive inputs — gradients stay 1 (or 0) regardless of magnitude. This avoids the vanishing-gradient problem of older activations like sigmoid or tanh.
• Sparsity: ~half the units are zero on average, which has nice computational and statistical properties.

Could you use Mish, GELU, SiLU instead? Sure. ReLU is the textbook default and works fine here.

Why sigmoid on the output

The action is a target y-position, normalized to [0, 1]. The sigmoid activation:

maps any real input to (0, 1). This is the natural activation for outputs that need to be in [0, 1]. (Compare to HW2 P2's actor, which used tanh for the [-1, 1] action range.)

The full architecture in PyTorch

Just nn.Sequential stacking. Five operations, each one line:

nn.Sequential(
    nn.Linear(state_dim, hidden),    # 4 → 256
    nn.ReLU(),
    nn.Linear(hidden, hidden),       # 256 → 256
    nn.ReLU(),
    nn.Linear(hidden, action_dim),   # 256 → 20
    nn.Sigmoid(),
)

That's the entire model. We'll write this in your BCPolicy.__init__.

The MSE Loss

The math

Read it: "average over the dataset of the squared distance between the policy's predicted action and the expert's action."

For the action chunk version, the squared norm is over all 20 dimensions of the chunk:

Each future action's prediction error squared, summed.

Why squared error and not absolute error?

For BC on smooth control tasks, MSE is the standard choice. The minimizer is the conditional expectation of expert actions:

This is a fact you should remember. The MSE-minimizing prediction at state s is the average of all the expert's actions at that state.

For unimodal experts (always make the same choice), this is exactly what we want. For multimodal experts (sometimes choose A, sometimes B), the mean of A and B is a problem — it's neither A nor B, and it could be a terrible prediction. That's the next chapter.

The Multimodality Trap

This chapter is the conceptual climax of Problem 1. Spend time here.

The setup

Hard mode has alternating single-gap and double-gap pipes. When the expert sees a double-gap pipe, it commits to either gap 1 or gap 2 randomly — from the same observation. So the dataset contains:

• State s (some double-gap configuration), action = aim at gap 1 (top, say y=0.7).
• State s (same exact state), action = aim at gap 2 (bottom, say y=0.3).

Two valid expert actions, randomly chosen, from the same state.

What MSE does

The MSE-minimizer is the conditional mean. Average of 0.7 and 0.3 is 0.5. So the trained policy outputs 0.5 at this state — which is exactly the wall between the two gaps.

The policy isn't doing anything wrong — it's faithfully reproducing the conditional mean of the expert distribution. But the conditional mean is not a valid action.

Why this is the deep lesson of HW1

Problem 1 is designed to fail on hard mode. You'll observe poor performance — mean episode length much shorter than 1000 — and the homework asks you to explain why in 2-3 sentences. The answer is exactly: multimodality of the expert plus mean-seeking behavior of MSE.

The two follow-up problems are direct fixes:

• Problem 2 (Flow Matching): train a generative model that can capture multimodal distributions. At each state, the model can output either gap 1 or gap 2, not the average.
• Problem 3 (DAgger): introduce a deterministic expert that picks one mode consistently, giving the BC policy clean unimodal supervision.

Both are valid solutions to the same problem. They take orthogonal approaches: Problem 2 makes the model richer; Problem 3 changes the data to be unimodal. Real robot learning systems often combine both: rich generative models + curated expert data.

Your one-paragraph writeup template

The PDF asks for 2-3 sentences explaining MSE on hard mode. Here's the structure:

PyTorch Primer

If you know Python but PyTorch is new, read this. You'll only need a small slice of PyTorch for Problem 1.

Tensors

Numpy arrays with two extras: GPU residency and autograd. torch.tensor([1., 2., 3.]) creates a tensor; .shape, .mean(), +, *, etc. work like numpy.

nn.Module: a network

Inherit from nn.Module, define layers in __init__, define the forward pass in forward. The framework handles parameters and gradient tracking automatically.

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(10, 256)
        self.layer2 = nn.Linear(256, 1)

    def forward(self, x):
        x = torch.relu(self.layer1(x))
        return self.layer2(x)

net = SimpleNet()
y = net(torch.randn(32, 10))    # shape [32, 1]

Notice net(x) calls forward(x) indirectly. The framework runs some bookkeeping then forwards. Don't call .forward() directly — you'll skip the bookkeeping and break things.

nn.Sequential: stacking layers

For pure feed-forward networks (no skip connections, no branching), you can skip the boilerplate of __init__+forward and just stack layers:

self.net = nn.Sequential(
    nn.Linear(10, 256),
    nn.ReLU(),
    nn.Linear(256, 1),
)

y = self.net(x)        # runs through all layers in order

Same effect as the manual class. Internally, nn.Sequential just calls each layer in order. You'll use this pattern in BCPolicy.

Loss functions

PyTorch's loss functions are usually classes you instantiate, then call:

criterion = nn.MSELoss()
loss = criterion(predicted, target)        # scalar

Or use the functional form, which is one-shot:

import torch.nn.functional as F

loss = F.mse_loss(predicted, target)        # same thing, no instance needed

Both apply .mean() by default — reducing the per-element squared errors to a scalar by averaging. The default is what you want for SGD.

Shapes

For BC on this homework:

• Input state has shape [B, 4] — B = batch size (e.g., 64), 4 = state dimension.
• Output predicted_action has shape [B, 20] — B batches, 20-dim action chunks.
• Target expert_action has shape [B, 20] — same.

The MSE loss reduces both to a scalar by squaring elementwise then averaging across all 20 × B elements.

Activation functions

You'll use nn.ReLU() twice and nn.Sigmoid() once in BCPolicy.

That's all the PyTorch you need for Problem 1. Onward.

Code Tour

Three files for Problem 1.

BCPolicy skeleton

Look at networks.py:316-339. The class is empty:

class BCPolicy(nn.Module):
    def __init__(self, state_dim: int = 4, action_dim: int = 20, hidden: int = 256):
        super().__init__()
        # TODO: Implement BCPolicy.__init__
        raise NotImplementedError("TODO: Implement BCPolicy.__init__")

    def forward(self, state):
        # TODO: Implement BCPolicy.forward
        raise NotImplementedError("TODO: Implement BCPolicy.forward")

Default constructor args: state_dim=4, action_dim=20, hidden=256. These match the environment's observation size, the action chunk length, and the standard hidden width.

mse_loss skeleton

Look at losses.py:18-33. Same pattern:

def mse_loss(policy, s_batch, a_batch):
    # TODO: Implement mse_loss
    raise NotImplementedError("TODO: Implement mse_loss")

The function signature receives policy (a callable that takes states and returns predicted actions), s_batch of shape [B, state_dim], and a_batch of shape [B, action_dim]. Returns a scalar.

The training loop (already done)

You don't write this, but it's worth understanding what calls your code. From main.py, paraphrased:

# Once at the start: collect demos, build dataset
states, actions = collect_expert_data(...)
dataset = TensorDataset(states, actions)
loader = DataLoader(dataset, batch_size=64, shuffle=True)

# Build policy and optimizer
policy = BCPolicy()
optimizer = optim.Adam(policy.parameters(), lr=1e-3)

# Standard supervised learning loop
for epoch in range(num_epochs):
    for s_batch, a_batch in loader:
        loss = mse_loss(policy, s_batch, a_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Plain SGD on a regression task. Your mse_loss goes inside the inner loop. Your BCPolicy is what policy(s_batch) calls.

What's nice about this setup

The training loop is generic. To switch from MSE regression (Problem 1) to flow matching (Problem 2), you just swap the loss function and policy class. Same loop, same dataset, same optimizer. main.py dispatches on the --method CLI flag and picks the right combination. This is the same pattern HW2/HW3 used — the framework is generic, your code is the algorithm.

Your Three Changes, Decoded

Per-line annotations for every blank you'll fill in. This is the centerpiece chapter.

def __init__(self, state_dim=4, action_dim=20, hidden=256):
    super().__init__()
    self.net = nn.Sequential(
        nn.Linear(state_dim, hidden),
        nn.ReLU(),
        nn.Linear(hidden, hidden),
        nn.ReLU(),
        nn.Linear(hidden, action_dim),
        nn.Sigmoid(),
    )

self.net = nn.Sequential(...) # without sigmoid
def forward(self, x):
    return torch.sigmoid(self.net(x))

def forward(self, state):
    return self.net(state)

self.fc1 = nn.Linear(state_dim, hidden)
self.fc2 = nn.Linear(hidden, hidden)
self.fc3 = nn.Linear(hidden, action_dim)

def forward(self, state):
    x = torch.relu(self.fc1(state))
    x = torch.relu(self.fc2(x))
    return torch.sigmoid(self.fc3(x))

def mse_loss(policy, s_batch, a_batch):
    predicted = policy(s_batch)
    return nn.functional.mse_loss(predicted, a_batch)

predicted = policy(s_batch)
return ((predicted - a_batch) ** 2).mean()

Summary

Three blanks. Total < 15 lines of actual code. The conceptual content is the architecture choice (3-layer MLP) and the loss choice (MSE). Both are textbook. The interesting part is in Chapter 07: why MSE fails on hard mode.

Running It

Setup

From the starter code's installation.md:

conda create -n cs224r-hw1 python=3.10
conda activate cs224r-hw1
cd hw1_starter_code
pip install -e .

This installs the local hw1 package in editable mode plus its dependencies (torch, gymnasium, numpy, matplotlib).

Running BC regression on easy mode

python main.py --method bc_reg --env easy

This:

1. Collects expert demos on easy mode (~5 episodes, ~5000 transitions).
2. Trains your BCPolicy via mse_loss for some number of epochs.
3. Evaluates on 50 episodes, reports mean and std of episode length.
4. Saves results to bc_reg_easy.txt.

Expected on easy mode: mean episode length close to 1000 (full survival), low standard deviation. Easy mode has unimodal expert (always aim at the single gap), so MSE works fine.

Running BC regression on hard mode

python main.py --method bc_reg --env hard

Same training pipeline, hard mode. Expected: mean episode length much shorter than 1000, high standard deviation. The policy is the conditional mean of the multimodal expert and crashes into the wall on most double-gap pipes.

Save the result to bc_reg_hard.txt. Report both numbers in the writeup table.

How long does training take?

Probably 1-5 minutes on CPU, much faster on GPU. The model is small, the dataset is small, and there's no env interaction during training. This is one homework where you don't need a cluster.

What healthy training looks like

Common reasons it doesn't work

1. Forgot super().__init__(): the model has no parameters tracked, optimizer has nothing to update, training does nothing. Easy give-away: training loss never changes.
2. Forgot the Sigmoid: predictions are unbounded, env clamps but training is unstable. Random/poor performance.
3. Wrong action_dim: BCPolicy(action_dim=1) instead of 20 outputs single scalar instead of 20-step chunk. Loop crashes immediately.
4. Calling policy.forward(s_batch) instead of policy(s_batch): usually works for BC but bad style and can cause issues with newer PyTorch features.

The deliverable

Per the PDF, you need three things in your writeup:

1. Table for easy mode: mean ± std of episode length over 50 evaluation episodes.
2. Table for hard mode: same metric.
3. 2-3 sentence explanation of MSE's hard-mode performance. Use the template from Chapter 07.

Cheat Sheet & Self-Quiz

Equations

API reference

Self-quiz

1. What's the difference between imitation learning and reinforcement learning?
2. What is behavior cloning (in one sentence)?
3. Why does the BC policy use sigmoid on its output?
4. What does the MSE loss converge to in expectation?
5. Why does action chunking help BC?
6. What's the difference between predicting 20 actions and executing 20 actions?
7. Why does easy mode work but hard mode fail with MSE?
8. What does "the conditional mean of expert actions" mean for the gap-1-or-gap-2 case?
9. Why is super().__init__() required in BCPolicy.__init__?
10. What's the difference between policy(x) and policy.forward(x)?
11. Why use ReLU instead of sigmoid for hidden layers?
12. What two HW1 problems fix the multimodality issue we observe in P1, and how?

Implementation order

1. BCPolicy.__init__ — ~1 minute. Six lines, all standard nn.Module idioms.
2. BCPolicy.forward — 30 seconds. One line.
3. mse_loss — 1 minute. Two lines.

Total: ~3 minutes of typing. Run easy mode first, verify ~1000 step survival. Run hard mode, observe the failure, write the explanation. Move on to Problem 2 (flow matching) which fixes the multimodality issue.

You can now teach this

Three big ideas, in order of importance:

1. Behavior cloning is supervised learning on (state, action) pairs. There's nothing fancy about the algorithm. The interesting parts are: which loss function, which architecture, and what to do when supervision is multimodal.
2. The MSE-minimizer is the conditional mean. This is fine when the expert is deterministic per state. It's catastrophic when the expert is multimodal — the mean of two valid actions can be a forbidden middle ground (the wall in Flappy Bird's hard mode).
3. Two orthogonal fixes for multimodality: enrich the model (Problem 2's flow matching can capture multimodal distributions) or sanitize the data (Problem 3's DAgger uses a deterministic expert). Both work; both are widely used in real robot learning.

If a friend asks: "Why does behavior cloning fail on hard mode?" — you say: "It's not BC's fault — it's the loss function. MSE regression converges to the conditional mean of the expert's actions. When the expert is bimodal — sometimes go up, sometimes go down — the mean is in the middle, which is a wall. The fix is either a richer model (flow matching captures bimodality) or cleaner data (a deterministic expert that always picks the same gap)."

You can teach this. On to Problem 2.