🤖 Robot Delivery Academy

📜The task, in plain English

A delivery robot lives on an 8×8 grid city. Each episode: start somewhere, drive to the depot holding a package, pickup, drive to the destination depot, dropoff. Walls block movement, and every map is a bit different. You must train a model that learns this behavior purely from expert demonstrations — supervised pairs of observation → action. Writing a path-finding algorithm (A*, BFS, planning) is explicitly forbidden: they're testing whether a model can learn the behavior, not whether you can code a solver.

🔧The baseline you're given

A complete, honest baseline is provided — and its own description lists its flaws:

1. Dataset: flattens every observation — the 6×8×8 grid is reshaped into a 384-long vector and concatenated with the 13 features → 397 numbers.
2. Model: a two-hidden-layer MLP (397 → 128 → 128 → 6), cross-entropy loss, Adam, 30 epochs.
3. Inference: picks the argmax action; the action_mask is applied only at inference time, not during training.
4. Evaluation: full-episode rollouts in the provided simulator, with GIF replays of episodes so you can literally watch it fail.

Stated limitations (= your to-do list, verbatim from the notebook): flattening destroys spatial structure; rare actions (pickup/dropoff occur ~2× per episode vs dozens of moves) are under-learned; the mask isn't used in training; the architecture is deliberately small.

🚀Baseline vs. solution

🧒Explain it like I'm brand new

Analogy: learning to drive by watching dashcam videos of a perfect driver. You learn "in this exact situation, she turned left." Two problems: (1) if you only ever memorize photo→action pairs, you never learn the concept of a road; (2) the moment you make one mistake she never made — you drift onto the shoulder — you're in a situation that appeared in zero videos, and you panic. That's behavioral cloning and its compounding-error problem in a nutshell.

• Why a CNN? An image (or a grid) has local structure — nearby cells matter to each other. A CNN is a model that looks through a small sliding window, so it learns "wall next to me" once, instead of 64 times.
• Why do rare actions need help? If 95% of your flashcards say "drive forward," you'll answer "drive forward" to everything and still score 95% — while failing every pickup. Weighting tells the loss: the rare cards count more.
• What's an action mask? A list of which moves are legal right now. Using it during training is like removing impossible answers from a multiple-choice exam before studying.

💬The Gemma 4 playthrough (2000-token limit)

Beginner here. This is the Dataset + model from a baseline I must improve (behavioral cloning, 8x8 grid, 6 actions). Explain in 8 bullet points max what it does and its 3 biggest weaknesses:
<paste the ~40 lines of Dataset + MLP code>

• Collects (observation, action) pairs from every expert trajectory step
• Flattens the 6×8×8 grid + 13-dim vector into one 397-dim input
• 2-layer MLP classifies the next action, cross-entropy loss
• Action mask applied only at inference
Weaknesses:
1. Flattening discards spatial structure a CNN would exploit
2. pickup/dropoff are rare classes — will be under-predicted
3. Per-action accuracy is optimized but episode success is scored

Write ONLY a PyTorch nn.Module: input a (B,6,8,8) grid tensor and a (B,13) vector. Two Conv2d layers (6→32→64, 3x3, padding=1, ReLU), flatten, concat the vector, then Linear(64*64+13 → 128) → ReLU → Linear(128 → 6). Code only, max 35 lines.

import torch, torch.nn as nn

class GridPolicy(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(6, 32, 3, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(),
        )
        self.head = nn.Sequential(
            nn.Linear(64*8*8 + 13, 128), nn.ReLU(),
            nn.Linear(128, 6),
        )
    def forward(self, grid, vec):
        z = self.conv(grid).flatten(1)
        return self.head(torch.cat([z, vec], dim=1))

My action labels are ints 0-5 in a list y_all. Actions 4 and 5 are rare. Give ONLY code to compute inverse-frequency class weights and use them in nn.CrossEntropyLoss. Max 10 lines.

Each sample has action_mask, a (6,) 0/1 array of legal actions. Modify this line of my training step so illegal logits become -1e9 before the loss:
loss = criterion(model(grid, vec), y)
Give only the replacement lines.

logits = model(grid, vec)
logits = logits.masked_fill(mask == 0, -1e9)
loss = criterion(logits, y)

My per-action accuracy is 0.96 but episode success rate is only 0.55. Replays show the robot reaches the package depot but oscillates north-south next to it without pressing pickup. 3 most likely causes + fastest fix for each, no code.

1. pickup under-predicted (rare class) → increase its class weight / oversample pickup steps
2. Model can't tell 'on depot' vs 'next to depot' → check the carrying/depot channels reach the conv input correctly
3. Deterministic argmax loops between two equal logits → add the action mask at inference and break ties, or sample from softmax with low temperature

🎯Takeaways & what Day 1 might do with this

• Pattern family: supervised classification wearing a robot costume — plus the imitation-learning lesson (distribution drift).
• The notebook's stated limitations are the answer key. CNN, rare-action weighting, mask-in-training: all three are literally listed for you.
• Diagnose by replaying episodes, not by staring at accuracy. The simulator + GIF tools are given to you for exactly this.
• Format discipline: predictions.jsonl inside predictions.zip, exact field names. Submit the plain baseline's zip in the first 30 minutes.
• Likely Day-1 extension: the notebook calls itself the "preparatory program for the Robot Training task" — expect bigger maps, partial observability, noisier demos, or multi-package episodes. The CNN + weighting + mask toolkit transfers directly.