HW1 · Imitation Learning

What you’re building

The assignment is imitation learning on Push-T: a 5-dimensional state describes the positions of a T-shaped block and the agent, the action is a 2-dimensional target position for the agent, and the goal is to push the T into a goal zone. You train two policies from the same expert dataset. First an MSE policy — a network that maps the current observation to an action chunk in a single forward pass, trained with mean-squared error. Then a flow matching policy, similar to diffusion policy, that generates chunks by integrating a learned vector field from noise.

Both are action chunking policies: at time $t$, the policy $\pi_\theta(\mathbf{A}_t \mid \mathbf{o}_t)$ predicts a whole chunk $\mathbf{A}_t = (\mathbf{a}_t, \mathbf{a}_{t+1}, \ldots, \mathbf{a}_{t+K-1})$ for a fixed length $K$, the environment executes it open-loop, and only then is the policy queried again on $\mathbf{o}_{t+K}$. Reduced decision frequency, one prediction per $K$ steps — the design used by modern robot imitation systems.

Part one: chunking with an MSE loss

The MSE policy is the simplest thing that could work: fit $\pi_\theta$ to expert observation–chunk pairs $(\mathbf{o}_t^{(j)}, \mathbf{A}_t^{(j)})$ by minimizing

$\mathcal{L}_{\text{MSE}}(\theta) = \frac{1}{B} \sum_{j=1}^{B} \left\| \mathbf{A}_t^{(j)} - \pi_\theta(\mathbf{o}_t^{(j)}) \right\|_2^2$

over batches of size $B$. Your work here is filling in the MSEPolicy TODO in src/hw1_imitation/model.py (the handout recommends a simple MLP with ReLU activations) and the training loop TODO in src/hw1_imitation/train.py. Before writing anything, read README.md, data.py (the dataset class — downloading and loading Push-T is handled for you), model.py (the BasePolicy class), and evaluation.py. You never modify evaluation.py, but you must call evaluate_policy periodically from your training loop — it logs metrics and rollout videos to WandB, and grading depends on those calls plus logger.dump_for_grading() at the end of training. The bar: an MSE policy should reach a reward of at least 0.5.

Part two: flow matching

An MSE policy predicts one chunk, so when the expert data is multimodal it regresses toward an average of the modes. Flow matching fixes this by learning a conditional vector field that transports noise into realistic chunks. Sample noise $\mathbf{A}_{t,0}^{(j)} \sim \mathcal{N}(0, I)$ and a timestep $\tau^{(j)} \sim \mathcal{U}(0,1)$, form the interpolation $\mathbf{A}_{t,\tau}^{(j)} = \tau^{(j)} \mathbf{A}_t^{(j)} + (1 - \tau^{(j)})\, \mathbf{A}_{t,0}^{(j)}$, and train $v_\theta$ to predict the velocity pointing from noise to data:

$\mathcal{L}_{\text{FM}}(\theta) = \frac{1}{B} \sum_{j=1}^{B} \left\| v_\theta(\mathbf{o}_t^{(j)}, \mathbf{A}_{t,\tau}^{(j)}, \tau^{(j)}) - (\mathbf{A}_t^{(j)} - \mathbf{A}_{t,0}^{(j)}) \right\|_2^2$

At inference, integrate the ODE $\frac{d\mathbf{A}_{t,\tau}}{d\tau} = v_\theta$ from $\tau = 0$ to $\tau = 1$ with $n$ Euler steps; $\mathbf{A}_{t,1}$ is the chunk you execute. Code-wise this part is only the FlowMatchingPolicy TODO in model.py — the handout suggests reusing the MSE policy’s MLP architecture, and warns (its one tip) not to forget that the network takes $\tau$ as an input. Target: reward of at least 0.7, notably above the MSE bar.

Which lectures feed it, and how to attack it

Lecture 2 is the theory backbone: behavioral cloning as supervised learning, why naive BC struggles (distribution shift, multimodal experts), and the remedies — this assignment exercises the expressive-policy-class branch (diffusion-style models) and the action-chunking idea rather than DAgger. Rewatch the multimodality segment if the flow matching loss feels unmotivated; Lecture 1 supplies the framing of policies as $\pi_\theta(\mathbf{a} \mid \mathbf{o})$. Order of attack: read the four starter files first, get the MSE training loop running end to end with evaluation calls wired in, and only then start the flow policy — its training loop is the one you already built. The conceptually new work is small; most debugging lives in tensor shapes (chunks are $K \times 2$, batches stack them) and in inference-time integration. Compute is not a bottleneck: a small MLP trains “within a few minutes” on a laptop CPU, Adam is the recommended optimizer, and torch.compile on the train step speeds things up significantly.

Deliverables and submission

For each policy you submit a working implementation, WandB logs (with videos) of a successful run, and a brief report: training curves of steps versus loss and reward, plotted yourself rather than screenshotted from WandB, plus a short description of your MLP architecture for part one and a qualitative comparison of how the flow policy behaves versus the MSE policy (based on the videos) for part two. Submission is a zip with an exp directory (your best run per part, renamed to mse and flow, keeping every .wandb, .json, .mp4, and .pkl file) and the src folder with the same structure as the original repo — the handout draws the exact tree; match it. Upload the zip to Gradescope under HW1 Code and the report under HW1 Report.