IOAI Training: Learn AI by solving problems

The problem in brief

A delivery robot lives on an $8\times 8$ map with 8 wall cells and 6 depots. In each episode it must drive to the depot holding a package, pick it up, drive to a second depot and drop it off. Every map is different. You do not write the route planner yourself: you get 400 recorded expert episodes (5 327 observation/action pairs on 100 different maps) and must train a model that imitates them, then lets it drive alone on 200 validation and 1 600 test scenarios on maps it has never seen. Planning or search may not be used to create extra demonstrations, and the final prediction process must be deterministic.

The data is public (the official Drive folder). Because the evaluator only replays your actions in the provided simulator, we can score the official test scenarios exactly by running that simulator ourselves.

How you are scored

The score is the success rate over the test scenarios:

\text{SR} = \frac{\#\{\text{episodes where the package is delivered within 120 steps}\}}{\#\{\text{episodes}\}}.

The statement publishes no threshold or reference score for this preparatory task, so there is no official target to compare against. A random policy scores about 0.

First look at the data

Three facts shape everything that follows.

The expert is a shortest-path planner. Every one of the 400 demonstrations has exactly the length of the BFS shortest route (start to package, package to destination, plus the two pick/drop actions). When several moves are equally short, the expert always prefers the lowest action id (south, then north, then east, then west).
The maps are new at test time. None of the 1 800 validation and test maps appears among the 100 training maps. The model has to learn how to navigate, not memorise routes.
The demonstrations only visit "good" states. The expert never makes a mistake, so the data contains no example of how to recover from a bad position. This is the classic weakness of behavioural cloning, and the statement warns about it: one wrong step can lead the robot into a state that the data never covered.

A first baseline

The notebook's baseline flattens the 6 grid planes and the 13-number vector into 397 inputs and trains a 2-layer MLP. It reaches 97 % action accuracy on the training pairs, but only 19.5 % success on validation: the high accuracy is memorisation of 100 maps. Simply masking the invalid actions (moves into walls, pick-up away from the package) during training and inference lifts it to 51 %. Still poor: a flattened grid gives the MLP no notion that cell $(3,4)$ is next to cell $(3,5)$ .

The key idea

Navigation is a local rule repeated everywhere on the map, so the model should be one too, and the remaining errors are loops, which the policy itself can break.

1. A convolutional, translation-equivariant policy. Which way to go depends on the walls around you and on where the target is, and this rule is the same in every corner of the map. A convolution applies the same small filter at every cell, so what it learns on one map transfers to every position of every other map. Five $3\times 3$ convolution layers give each cell a receptive field of $11\times 11$ , larger than the $8\times 8$ map: every cell can "see" the whole map. We then read out only the features of the robot's own cell and of its 4 neighbours (in the order south, north, east, west) and map them to the 6 action logits. The network effectively scores "how good is it to step onto this neighbour", which is the shape of the true answer (the neighbour's distance to the target).

We also simplified the input: instead of separate "package" and "destination" planes plus a carrying flag, one plane marks the cell to go to now (the package before pick-up, the destination after). The two halves of an episode become the same problem.

2. The 8 symmetries of the square as free data. Rotating or mirroring a map gives another valid map, and the correct move rotates with it (south becomes east after a quarter turn, and so on). Applying a random symmetry to each mini-batch, and permuting the action labels and the action mask accordingly, multiplies the 100 training maps by 8. Training accuracy drops to about 91 % because the expert's tie-break (prefer south) is not symmetric, but that does not hurt: any tie-breaking among shortest moves is equally good. Validation success goes from 80.5 % to 91 %.

3. Every remaining failure is a loop. We replayed the 18 failed validation episodes: all 18 end with the robot cycling between the same few cells until the 120-step limit, 12 of them before the pick-up and 6 after. The policy is deterministic, so the moment it returns to a state $(\text{cell}, \text{carrying})$ it has already been in, it will repeat the same cycle forever. The figure shows a typical case: the robot steps east into a pocket, decides west is best, then east again.

A loop caused by a deterministic policy, and the same policy with loop breaking

The fix uses only the model's own preferences: never take the same action twice from the same state. If the best action from this state was already tried, it led back here, so it is part of a cycle; take the next-best action instead. Formally, with $T(s)$ the set of actions already taken in state $s$ ,

a_t = \arg\max_{a \,\text{valid},\ a \notin T(s_t)} \ \text{logit}_\theta(s_t)_a .

This is still deterministic and uses no planner and no knowledge of the map beyond what the model outputs. Each state has at most 4 moves, so the robot cannot cycle forever: it is forced to explore the second choices of the model, which are usually right when the first choice was wrong.

The full solution

Convert the data safely. The three .pkl files come from a third party; we opened them once in a sandbox (no network, restricted writes) and saved scenarios and action lists as JSON. Observations are rebuilt by replaying the actions in a re-implementation of the simulator.
Build the 5 input planes for each of the 5 327 states: wall, free cell, current target, robot, carrying.
Train the policy (64 channels, 5 conv layers, readout of 5 cells, about 150 k parameters): 40 epochs, Adam with a one-cycle learning rate up to $2\cdot 10^{-3}$ , batch 128, a random dihedral transform per batch, cross-entropy on masked logits. About 2 minutes on 2 CPU threads.

k = rng.integers(8); p = torch.tensor(action_perm(k))   # where each action goes under symmetry k
x, y = dihedral(X[b], k), p[Y[b]]
m = torch.zeros_like(M[b]); m[:, p] = M[b]               # the mask rotates with the actions
loss = F.cross_entropy(model(x).masked_fill(~m, -1e9), y)

Roll out each test scenario with the masked argmax and loop breaking, and write the action lists to predictions.jsonl inside predictions.zip.

Things that did not help: a value-iteration style network (one convolution applied 20 times with shared weights) fitted the training pairs perfectly (100 % accuracy) but generalised slightly worse (88.5 % without augmentation, 91 % with), and averaging the logits over the 8 symmetries at test time or ensembling two models changed little once loops were handled.

Results

Success rate in the official simulator. Validation = the 200 provided validation scenarios; test = the 1 600 official test scenarios (scored by replaying our actions, exactly like the evaluator).

approach	valid SR	test SR
random actions	about 0
baseline MLP (flattened input)	0.195
MLP + action masking	0.510
CNN, 5-cell readout, masking	0.805
value-iteration style recurrent CNN	0.885
recurrent CNN + symmetry augmentation	0.910
CNN + symmetry augmentation	0.915	0.942
+ test-time averaging over 8 symmetries	0.915
CNN + augmentation + loop breaking (final)	0.990	0.992

The final model delivers the package in 1 587 of the 1 600 test episodes (99.2 %). Loop breaking alone rescues 80 test episodes. The routes are also efficient: 14.2 steps on average against 13.1 for the shortest possible routes, and 92.5 % of the successful episodes are exactly as short as the expert's. An earlier training run of the same architecture gave 0.995 on validation, so the seed-to-seed variation is about half a point.

What to remember

Accuracy on demonstrations is not success. 97 % action accuracy gave 19.5 % success. In sequential problems, evaluate the full rollout, and look at how episodes fail.
Match the architecture to the structure. A grid task with a rule that is the same everywhere calls for convolutions, and a readout at the agent's position turns "classify the whole map" into "score each neighbour".
Symmetries are free data. Rotations and mirrors of a map are new valid maps, provided you transform the labels (and the masks) consistently.
A deterministic policy that revisits a state is trapped. Detecting a repeated state and falling back on the model's next-best action is cheap, deterministic, and fixes the covariate-shift failures that behavioural cloning cannot learn to avoid from perfect demonstrations alone.

"""Robot Delivery Academy (IOAI 2026 at-home task 2): behavioural cloning on 8x8 delivery maps.

Pipeline: rebuild (observation, expert action) pairs from the 400 demonstrations, train a small
fully-convolutional policy with the 8 symmetries of the square as data augmentation, then roll it
out in the simulator with a "never repeat the same action from the same state" rule that breaks
the loops a deterministic policy can fall into. Writes predictions.zip for the 1600 test scenarios
and reports the success rate (the simulator IS the official evaluator: replay of the actions).

Data: gdown folder 1DXFDoY9bqulMBFacyDShVIx8Sa7Z5Wpa -> data/*.pkl. The pickles come from a third
party, so convert them ONCE inside a sandbox (no network, restricted writes):
    run_sandboxed.sh . python3 solution.py --convert      (writes train.json, valid.json, test.json)
then run normally:  python3 solution.py
"""
import collections, json, sys, zipfile
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F

MOVES = {0: (1, 0), 1: (-1, 0), 2: (0, 1), 3: (0, -1)}   # south, north, east, west; 4 pickup, 5 dropoff
MAX_STEPS = 120

def convert():
    import pickle
    def js(o):
        if isinstance(o, dict): return {k: js(v) for k, v in o.items()}
        if isinstance(o, (list, tuple)): return [js(v) for v in o]
        if isinstance(o, np.ndarray): return o.tolist()
        return o.item() if isinstance(o, np.generic) else o
    tr = pickle.load(open("data/train_demos.pkl", "rb"))["trajectories"]
    json.dump([{"scenario": js(t["scenario"]), "actions": js(t["actions"])} for t in tr], open("train.json", "w"))
    for n in ("valid", "test"):
        json.dump(js(pickle.load(open(f"data/{n}_scenarios.pkl", "rb"))), open(f"{n}.json", "w"))

class Sim:
    """Same rules as the notebook's DeliverySimulator8x8."""
    def __init__(self, s):
        self.walls = {tuple(c) for c in s["walls"]}; self.depots = [tuple(c) for c in s["depots"]]
        self.pos = tuple(s["agent_pos"]); self.pkg = s["package_location"]; self.dst = s["destination"]
        self.carrying = False; self.t = 0
    def free(self, r, c): return 0 <= r < 8 and 0 <= c < 8 and (r, c) not in self.walls
    def target(self): return self.depots[self.dst if self.carrying else self.pkg]
    def mask(self):
        m = np.zeros(6, bool); r, c = self.pos
        for a, (dr, dc) in MOVES.items(): m[a] = self.free(r + dr, c + dc)
        m[4] = not self.carrying and self.pos == self.depots[self.pkg]
        m[5] = self.carrying and self.pos == self.depots[self.dst]
        return m
    def step(self, a):
        done = False
        if a in MOVES:
            r, c = self.pos[0] + MOVES[a][0], self.pos[1] + MOVES[a][1]
            if self.free(r, c): self.pos = (r, c)
        elif a == 4 and self.mask()[4]: self.carrying = True
        elif a == 5 and self.mask()[5]: done = True
        self.t += 1
        return done, self.t >= MAX_STEPS
    def planes(self):
        """5 input planes: wall, free cell, CURRENT target (package or destination), robot, carrying.
        Merging package/destination into one 'where to go now' plane is the main representation choice."""
        x = np.zeros((5, 8, 8), np.float32)
        for w in self.walls: x[0][w] = 1
        x[1] = 1 - x[0]; x[2][self.target()] = 1; x[3][self.pos] = 1; x[4] = float(self.carrying)
        return x

class Policy(nn.Module):
    """5 conv layers (receptive field 11 > 8: every cell sees the whole map), then read the features
    of the robot's cell and its 4 neighbours (in action order) and map them to 6 logits."""
    def __init__(self, h=64, layers=5):
        super().__init__()
        self.convs = nn.ModuleList([nn.Conv2d(5 if i == 0 else h, h, 3, padding=1) for i in range(layers)])
        self.out = nn.Linear(5 * h, 6)
    def forward(self, x):
        h = x
        for conv in self.convs: h = F.relu(conv(h))
        hp, ap = F.pad(h, (1, 1, 1, 1)), F.pad(x[:, 3:4], (1, 1, 1, 1))
        cells = [(hp * torch.roll(ap, d, (2, 3))).sum((2, 3)) for d in [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)]]
        return self.out(torch.cat(cells, 1))

def dihedral(x, k):       # one of the 8 symmetries of the square
    if k & 1: x = x.flip(2)
    if k & 2: x = x.flip(3)
    if k & 4: x = x.transpose(2, 3)
    return x

def action_perm(k):       # where each action goes under symmetry k
    perm = []
    for a in range(4):
        dr, dc = MOVES[a]
        if k & 1: dr = -dr
        if k & 2: dc = -dc
        if k & 4: dr, dc = dc, dr
        perm.append([b for b, d in MOVES.items() if d == (dr, dc)][0])
    return perm + [4, 5]

def build_dataset(demos):
    X, M, Y = [], [], []
    for d in demos:
        s = Sim(d["scenario"])
        for a in d["actions"]:
            X.append(s.planes()); M.append(s.mask()); Y.append(a); s.step(a)
    return torch.tensor(np.array(X)), torch.tensor(np.array(M)), torch.tensor(Y)

def train(X, M, Y, epochs=40, bs=128, lr=2e-3, seed=0):
    torch.manual_seed(seed); rng = np.random.default_rng(seed)
    model = Policy(); opt = torch.optim.Adam(model.parameters(), lr=lr)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, lr, total_steps=epochs * ((len(Y) + bs - 1) // bs))
    for _ in range(epochs):
        idx = torch.randperm(len(Y))
        for i in range(0, len(Y), bs):
            b = idx[i:i + bs]; k = int(rng.integers(8)); p = torch.tensor(action_perm(k))
            x, y = dihedral(X[b], k), p[Y[b]]
            m = torch.zeros_like(M[b]); m[:, p] = M[b]          # move the mask with the actions
            logits = model(x).masked_fill(~m, -1e9)             # invalid actions never compete
            loss = F.cross_entropy(logits, y)
            opt.zero_grad(); loss.backward(); opt.step(); sched.step()
    return model.eval()

@torch.no_grad()
def rollout(model, scenario, unloop=True):
    s = Sim(scenario); actions = []; tried = collections.defaultdict(set)
    while True:
        logits = model(torch.tensor(s.planes())[None])[0].masked_fill(~torch.tensor(s.mask()), -1e9)
        key = (s.pos, s.carrying)
        if unloop:                     # an action already taken from this exact state led back here:
            for b in tried[key]:       # it is part of a loop, so fall back on the next-best action
                logits[b] -= 1e6
        a = int(logits.argmax()); tried[key].add(a); actions.append(a)
        done, timeout = s.step(a)
        if done or timeout: return done, actions

def evaluate(model, scenarios, unloop=True):
    res = [rollout(model, s, unloop) for s in scenarios]
    return np.mean([r[0] for r in res]), np.mean([len(r[1]) for r in res]), res

if __name__ == "__main__":
    if "--convert" in sys.argv: convert(); sys.exit()
    torch.set_num_threads(2)
    demos = json.load(open("train.json")); valid = json.load(open("valid.json")); test = json.load(open("test.json"))
    X, M, Y = build_dataset(demos)
    print("samples", len(Y))
    model = train(X, M, Y)
    for name, sc in [("valid", valid), ("test", test)]:
        for ul in (False, True):
            sr, steps, _ = evaluate(model, sc, ul)
            print(f"{name:5s} unloop={ul}: success rate {sr:.4f}, avg steps {steps:.1f}")
    _, _, res = evaluate(model, test, True)
    with open("predictions.jsonl", "w") as f:
        for s, (_, acts) in zip(test, res):
            f.write(json.dumps({"layout_id": s["layout_id"], "episode_seed": s["episode_seed"], "actions": acts}) + "\n")
    with zipfile.ZipFile("predictions.zip", "w") as z: z.write("predictions.jsonl")