ExplodingCB2 commited on
Commit
ee68aa6
·
verified ·
1 Parent(s): daaedeb

drive-v5 and the checkpoints it came from, with a model card

Browse files
README.md ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: pytorch
3
+ pipeline_tag: image-segmentation
4
+ tags:
5
+ - autonomous-driving
6
+ - go-kart
7
+ - racing-line
8
+ - semantic-segmentation
9
+ - imitation-learning
10
+ - robotics
11
+ ---
12
+
13
+ # McQueen driving model (drive-v5)
14
+
15
+ A small camera-only driving model for an autonomous go-kart on the Purdue Grand Prix track. One RGB frame goes in; three things come out: a **pavement mask**, a **racing line** (12 waypoints, 2 to 24 m ahead, in the kart's frame) and a **target speed** at each waypoint. It does not output steering or throttle. Controllers outside the model turn the line into commands from measured state, and a safety layer keeps the line inside the pavement the camera can see.
16
+
17
+ > **Status: a simulator result, not a road-tested model.** It has never seen the kart's own camera and has never driven a real kart. See [What this does not show](#what-this-does-not-show).
18
+
19
+ Code, training scripts, the simulator and the full version history are in the project repository: https://github.com/ExplodingCB/mcqueen-vip (branch `training-perception`, see `docs/13-driving-model-versions.md`).
20
+
21
+ ## Files
22
+
23
+ | File | What it is |
24
+ | --- | --- |
25
+ | `drive-v5.pt` | **The recommended model.** Real-footage mask training and a clean simulated lap |
26
+ | `checkpoints/drive-v3.pt` | Best simulated lap time (63.8 s); weaker on real footage |
27
+ | `checkpoints/drive-v4.pt` | drive-v3 plus real helmet-cam footage for the mask |
28
+ | `checkpoints/drive-v6.pt` | Trained on a faster teacher line. Best open-loop numbers, but **does not complete laps**; included as a research result |
29
+ | `model.py`, `driver.py`, `example.py` | Everything needed to load a checkpoint and run it on an image |
30
+ | `laptime-v2-best.json` | The searched expert line (racing line offsets and speed limits) the faster teacher uses |
31
+
32
+ ## Checksums (SHA-256)
33
+
34
+ | File | SHA-256 | Size |
35
+ | --- | --- | --- |
36
+ | `checkpoints/drive-v3.pt` | `ce23623f702185460a999e423af3ac36670f49dfd3e08d41c65285bb1132ac5c` | 15.7 MB |
37
+ | `checkpoints/drive-v4.pt` | `3d995f317eb3467ac0ca801500d060651d1ea6b3f68cea60bd9050f400431677` | 15.7 MB |
38
+ | `checkpoints/drive-v6.pt` | `b1a7813b21b4e79b55e7c2f1c3f9a244fb7641a0550674d96e9f4e3058e057b5` | 15.7 MB |
39
+ | `drive-v5.pt` | `0d6ba0f1a4e45fd28c4f609fbfb0cb7256a9feaf2071bab200498a0fcc12ddcc` | 15.7 MB |
40
+
41
+ ## Use
42
+
43
+ ```python
44
+ # pip install torch numpy pillow
45
+ import numpy as np
46
+ from PIL import Image
47
+ from driver import Driver # driver.py and model.py sit beside it
48
+
49
+ driver = Driver("drive-v5.pt") # CPU by default
50
+ frame = np.asarray(Image.open("frame.jpg").convert("RGB").resize((480, 270))) # uint8, 270 x 480 x 3
51
+ mask, line, speed = driver.predict_drive(frame)
52
+ # mask : (270, 480) pavement probability
53
+ # line : (12, 2) metres in the kart frame, x forward and y left, from 2 m to 24 m ahead
54
+ # speed: (12,) metres per second the model wants at each waypoint
55
+ ```
56
+
57
+ The model assumes a camera about 0.7 m above the ground, pitched about 8 degrees down, with a 100 degree horizontal field of view. A different mount shifts the line, because it is in metres. `python example.py frame.jpg` prints the same thing.
58
+
59
+ ## Model
60
+
61
+ `DriveNet`: a 3.9 M parameter encoder-decoder (a dilated bottleneck, all standard convolutions) with a mask head, plus a driving head that reads the bottleneck and the quarter-resolution decoder features (with pixel-coordinate channels) and regresses the line and speeds. Input is the 240 x 135 half-resolution frame. About 7 ms per frame on one CPU thread.
62
+
63
+ ## Training data
64
+
65
+ * **Simulator frames:** the Purdue track with the 2023 Indiana orthophoto draped over a flat ground plane (Indiana's imagery service, marked CC0), camera mount and image appearance randomized, labeled with an expert's minimum-curvature racing line and the speed along it.
66
+ * **DAgger frames:** frames from states the model itself reaches in closed loop (with steering bumps), labeled by the expert from where the kart actually is.
67
+ * **Real kart footage, for the mask only:** about 2,300 frames from about 30 YouTube kart onboard videos (helmet cams and nose or bumper cams), pavement labeled automatically by SegFormer (ADE20K classes road, runway, sidewalk, path). **These are pseudo-labels with no human review.** The videos are third-party, so check their terms before redistributing or using this commercially.
68
+
69
+ ## Results (all in the simulator unless stated)
70
+
71
+ | Measure | drive-v5 |
72
+ | --- | --- |
73
+ | Closed-loop lap, camera only, default limits, no speed cap, 3 seeds | **64.6 s**, 0.39 m closest to an edge |
74
+ | Racing-line error on unseen simulator frames, 2 to 10 m / 12 to 24 m ahead | 0.27 m / 0.54 m |
75
+ | Target-speed error | 0.34 m/s |
76
+ | Pavement mask IoU on unseen simulator frames | 0.979 |
77
+ | **Pavement mask IoU on 575 held-out real kart frames, against auto-labels** | **0.84** |
78
+
79
+ For reference: a simulator-only model scores 0.39 on the real frames, below a fixed "the lower part of the picture is road" guess (0.47). Models trained on the faster teacher (drive-v6) fail to complete laps.
80
+
81
+ ## What this does not show
82
+
83
+ * No version has seen the kart's own camera or driven a real kart. The held-out real frames are other people's cameras on other tracks, scored against automatic labels.
84
+ * The racing line and speed were trained only on simulator frames. The mask is the only output that has been trained on real images, so how the line and speed behave on real frames is untested.
85
+ * The track file is hand-traced from aerial imagery, and the simulated tyres, brakes and steering servo are placeholders. A 64.6 s simulated lap is not a prediction of the real kart.
86
+ * The driving path (safety filter, corridor tracker) exists only in the Python simulator and is not yet ported to the kart's software or timed on the Jetson.
87
+
88
+ The intended use today is research, the simulator, and **shadow mode** (a person drives while the model only watches).
checkpoints/drive-v3.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ce23623f702185460a999e423af3ac36670f49dfd3e08d41c65285bb1132ac5c
3
+ size 15691971
checkpoints/drive-v4.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3d995f317eb3467ac0ca801500d060651d1ea6b3f68cea60bd9050f400431677
3
+ size 15691907
checkpoints/drive-v6.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b1a7813b21b4e79b55e7c2f1c3f9a244fb7641a0550674d96e9f4e3058e057b5
3
+ size 15691907
drive-v5.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0d6ba0f1a4e45fd28c4f609fbfb0cb7256a9feaf2071bab200498a0fcc12ddcc
3
+ size 15691907
driver.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """The driving model behind the simulator's camera-model contract.
2
+
3
+ python -m mcq_sim evaluate --policy stack --perception driver:create
4
+
5
+ with ``PYTHONPATH=training/perception:src/mcq_sim``. Weights come from
6
+ ``MCQ_DRIVER_WEIGHTS``, else ``runs/drive-latest/model.pt`` beside this file.
7
+
8
+ ``predict(rgb)`` is the pavement probability, as for a segmenter, so the stack
9
+ can build its corridor. ``predict_drive(rgb)`` adds the racing line the model
10
+ intends to drive, as waypoints in base_link (x forward, y left), and the speed
11
+ it wants at each. The stack drives that line when it stays inside the corridor
12
+ and falls back to its own planner when it does not.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import os
18
+ from pathlib import Path
19
+
20
+ import numpy as np
21
+ import torch
22
+ import torch.nn.functional as F
23
+
24
+ from model import DriveNet
25
+
26
+ HERE = Path(__file__).resolve().parent
27
+
28
+
29
+ class Driver:
30
+ def __init__(self, weights: str | Path, device="cpu"):
31
+ torch.set_num_threads(1) # many drivers run side by side in rollouts; one thread each
32
+ checkpoint = torch.load(weights, map_location=device, weights_only=True)
33
+ self.net = DriveNet(fine=bool(checkpoint.get("fine", False))).to(device).eval()
34
+ self.net.load_state_dict(checkpoint["state"])
35
+ self.device = device
36
+ self.weights = str(weights)
37
+
38
+ @torch.no_grad()
39
+ def predict_drive(self, rgb: np.ndarray):
40
+ """uint8 HWC RGB (270 x 480) in; (pavement probability 270 x 480,
41
+ waypoints K x 2 in metres, speeds K in m/s) out."""
42
+ x = torch.from_numpy(np.ascontiguousarray(rgb)).to(self.device).permute(2, 0, 1)[None].float()
43
+ h, w = x.shape[-2:]
44
+ logits, points, speeds = self.net.forward_all(F.avg_pool2d(x, 2))
45
+ probability = torch.sigmoid(F.interpolate(logits[None], size=(h, w), mode="bilinear", align_corners=False))
46
+ out = probability[0, 0].cpu().numpy(), points[0].cpu().numpy(), speeds[0].cpu().numpy()
47
+ return out
48
+
49
+ def predict(self, rgb: np.ndarray) -> np.ndarray:
50
+ return self.predict_drive(rgb)[0]
51
+
52
+
53
+ def create():
54
+ path = os.environ.get("MCQ_DRIVER_WEIGHTS", HERE / "runs" / "drive-latest" / "model.pt")
55
+ return Driver(path)
example.py ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Run a checkpoint on one image: python example.py frame.jpg [drive-v5.pt]"""
2
+ import sys
3
+
4
+ import numpy as np
5
+ from PIL import Image
6
+
7
+ from driver import Driver
8
+
9
+ image = sys.argv[1] if len(sys.argv) > 1 else None
10
+ weights = sys.argv[2] if len(sys.argv) > 2 else "drive-v5.pt"
11
+ driver = Driver(weights)
12
+ frame = (
13
+ np.asarray(Image.open(image).convert("RGB").resize((480, 270)))
14
+ if image
15
+ else np.full((270, 480, 3), 110, dtype=np.uint8)
16
+ )
17
+ mask, line, speed = driver.predict_drive(frame)
18
+ print(f"{weights}: pavement share {float((mask > 0.5).mean()):.2f}")
19
+ for k, ((x, y), v) in enumerate(zip(line, speed, strict=True)):
20
+ print(f" waypoint {k:2d}: {x:5.1f} m ahead, {y:+5.2f} m left, wants {v:4.1f} m/s")
laptime-v2-best.json ADDED
@@ -0,0 +1,78 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "theta": [
3
+ -0.020713883569653507,
4
+ 0.7986371171163591,
5
+ -0.06007090210526833,
6
+ -0.13692955258757095,
7
+ 0.3164787844956688,
8
+ 0.22910077232445988,
9
+ 0.4085683741777101,
10
+ 0.1732443452654597,
11
+ 0.31628061199629176,
12
+ 0.2936319802034716,
13
+ -0.007033094360031467,
14
+ 0.0013200730496496782,
15
+ 0.40243794379289366,
16
+ 0.765853695695947,
17
+ 0.5178593376332585,
18
+ 0.33366131041959013,
19
+ 1.0801067821758314,
20
+ 0.7981998198285803,
21
+ -0.3385555048128157,
22
+ -0.3040821312994424,
23
+ 1.6182991758783387,
24
+ 0.3389952005170975,
25
+ -0.22789996980313015,
26
+ 0.4039164824230546,
27
+ 0.4619468070190555,
28
+ 0.21125474388800541,
29
+ -0.5006066348041253,
30
+ -0.8610472716707411,
31
+ -0.20847536428974067,
32
+ -0.2394684425191484,
33
+ -0.9391657225767008,
34
+ -0.9217363080423312,
35
+ 0.06200212942970203,
36
+ 0.6460498410297608,
37
+ 0.9070343531284966,
38
+ 1.2993982884320447,
39
+ 0.6067228170245667,
40
+ 0.8575524174357774,
41
+ 0.9950192746873036,
42
+ -0.16623329240627815,
43
+ 4.464017710110649,
44
+ 1.7967631217917692,
45
+ 2.0052742170641062
46
+ ],
47
+ "score": 55.65999999998595,
48
+ "results": [
49
+ {
50
+ "passed": true,
51
+ "lap_s": 54.64999999998654,
52
+ "progress_m": 435.10000000004015,
53
+ "clearance_m": 0.2806985629754275,
54
+ "max_speed": 12.532168425742148,
55
+ "corrections": 14,
56
+ "stop": null
57
+ },
58
+ {
59
+ "passed": true,
60
+ "lap_s": 55.29999999998595,
61
+ "progress_m": 435.1000000000418,
62
+ "clearance_m": 0.2562416972740049,
63
+ "max_speed": 12.49933797767215,
64
+ "corrections": 12,
65
+ "stop": null
66
+ },
67
+ {
68
+ "passed": true,
69
+ "lap_s": 54.64999999998654,
70
+ "progress_m": 435.25000000004,
71
+ "clearance_m": 0.4983136894426503,
72
+ "max_speed": 12.505934752719071,
73
+ "corrections": 0,
74
+ "stop": null
75
+ }
76
+ ],
77
+ "generation": 103
78
+ }
model.py ADDED
@@ -0,0 +1,127 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """A small encoder-decoder for the pavement mask.
2
+
3
+ About 1.6 M parameters, all standard convolutions, so it exports to TensorRT
4
+ without surprises and runs far inside the 20 fps budget on the Orin. The input is
5
+ the 240 x 135 half-resolution frame; the output is one logit per pixel at the
6
+ same size. A dilated bottleneck gives the receptive field to see the whole road
7
+ shape at once, which matters more than texture: pavement and grass can be the
8
+ same color, and the shape of the corridor still says which is which.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import torch
14
+ import torch.nn as nn
15
+ import torch.nn.functional as F
16
+
17
+
18
+ def block(cin, cout, stride=1, dilation=1):
19
+ return nn.Sequential(
20
+ nn.Conv2d(cin, cout, 3, stride, padding=dilation, dilation=dilation, bias=False),
21
+ nn.BatchNorm2d(cout),
22
+ nn.ReLU(inplace=True),
23
+ nn.Conv2d(cout, cout, 3, 1, padding=dilation, dilation=dilation, bias=False),
24
+ nn.BatchNorm2d(cout),
25
+ nn.ReLU(inplace=True),
26
+ )
27
+
28
+
29
+ class PavementNet(nn.Module):
30
+ def __init__(self, width=24):
31
+ super().__init__()
32
+ w = width
33
+ self.e1 = block(3, w, stride=2) # 1/2
34
+ self.e2 = block(w, 2 * w, stride=2) # 1/4
35
+ self.e3 = block(2 * w, 4 * w, stride=2) # 1/8
36
+ self.e4 = block(4 * w, 6 * w, stride=2) # 1/16
37
+ self.context = nn.Sequential(block(6 * w, 6 * w, dilation=2), block(6 * w, 6 * w, dilation=4))
38
+ self.d3 = block(6 * w + 4 * w, 4 * w)
39
+ self.d2 = block(4 * w + 2 * w, 2 * w)
40
+ self.d1 = block(2 * w + w, w)
41
+ self.head = nn.Conv2d(w, 1, 1)
42
+
43
+ @staticmethod
44
+ def _up(x, like):
45
+ return F.interpolate(x, size=like.shape[-2:], mode="bilinear", align_corners=False)
46
+
47
+ def forward(self, x):
48
+ size = x.shape[-2:]
49
+ x = x / 255.0 - 0.5
50
+ e1 = self.e1(x)
51
+ e2 = self.e2(e1)
52
+ e3 = self.e3(e2)
53
+ e4 = self.context(self.e4(e3))
54
+ d3 = self.d3(torch.cat([self._up(e4, e3), e3], 1))
55
+ d2 = self.d2(torch.cat([self._up(d3, e2), e2], 1))
56
+ d1 = self.d1(torch.cat([self._up(d2, e1), e1], 1))
57
+ return F.interpolate(self.head(d1), size=size, mode="bilinear", align_corners=False)[:, 0]
58
+
59
+
60
+ class DriveNet(PavementNet):
61
+ """PavementNet with a driving head: from the same encoder, the racing line
62
+ as ``n_waypoints`` points in the kart's frame (x forward, y left) and the
63
+ target speed at each. The mask head stays, so the stack can still check the
64
+ driving output against the pavement it sees.
65
+
66
+ The head reads the whole bottleneck, not a pooled vector, because where the
67
+ road goes is a matter of where in the image it is. Outputs are scaled to
68
+ about unit size: x / 25 m, y / 10 m, speed / 15 m/s.
69
+ """
70
+
71
+ SCALE = (25.0, 10.0, 15.0)
72
+
73
+ def __init__(self, width=24, n_waypoints=12, grid=(9, 15), fine=False):
74
+ super().__init__(width)
75
+ self.n_waypoints = n_waypoints
76
+ self.squeeze = nn.Sequential(nn.Conv2d(6 * width, 32, 1), nn.BatchNorm2d(32), nn.ReLU(inplace=True))
77
+ self.grid = grid
78
+ # ``fine``: the head also reads the decoder's quarter-resolution features
79
+ # with two pixel-coordinate channels, then strides them down to the grid.
80
+ # Where the kart sits relative to the line is read from where the edges
81
+ # are 5 to 10 m ahead, a matter of a few pixels the 1/16 bottleneck blurs.
82
+ self.fine = fine
83
+ if fine:
84
+ self.fine_path = nn.Sequential(
85
+ nn.Conv2d(2 * width + 2, 64, 3, 2, 1, bias=False),
86
+ nn.BatchNorm2d(64),
87
+ nn.ReLU(inplace=True),
88
+ nn.Conv2d(64, 64, 3, 2, 1, bias=False),
89
+ nn.BatchNorm2d(64),
90
+ nn.ReLU(inplace=True),
91
+ nn.Conv2d(64, 32, 1),
92
+ nn.BatchNorm2d(32),
93
+ nn.ReLU(inplace=True),
94
+ )
95
+ head_in = 64 if fine else 32
96
+ self.drive = nn.Sequential(
97
+ nn.Flatten(),
98
+ nn.Linear(head_in * grid[0] * grid[1], 256),
99
+ nn.ReLU(inplace=True),
100
+ nn.Dropout(0.1),
101
+ nn.Linear(256, 3 * n_waypoints),
102
+ )
103
+
104
+ def forward_all(self, x):
105
+ """(mask logits N x H x W, waypoints N x K x 2 in metres, speeds N x K in m/s)."""
106
+ size = x.shape[-2:]
107
+ x = x / 255.0 - 0.5
108
+ e1 = self.e1(x)
109
+ e2 = self.e2(e1)
110
+ e3 = self.e3(e2)
111
+ e4 = self.context(self.e4(e3))
112
+ d3 = self.d3(torch.cat([self._up(e4, e3), e3], 1))
113
+ d2 = self.d2(torch.cat([self._up(d3, e2), e2], 1))
114
+ d1 = self.d1(torch.cat([self._up(d2, e1), e1], 1))
115
+ mask = F.interpolate(self.head(d1), size=size, mode="bilinear", align_corners=False)[:, 0]
116
+ code = F.adaptive_avg_pool2d(self.squeeze(e4), self.grid)
117
+ if self.fine:
118
+ n, _, h, w = d2.shape
119
+ ys = torch.linspace(-1, 1, h, device=d2.device).view(1, 1, h, 1).expand(n, 1, h, w)
120
+ xs = torch.linspace(-1, 1, w, device=d2.device).view(1, 1, 1, w).expand(n, 1, h, w)
121
+ detail = self.fine_path(torch.cat([d2, ys, xs], 1))
122
+ code = torch.cat([code, F.adaptive_avg_pool2d(detail, self.grid)], 1)
123
+ out = self.drive(code).view(-1, self.n_waypoints, 3)
124
+ sx, sy, sv = self.SCALE
125
+ points = torch.stack([out[..., 0] * sx, out[..., 1] * sy], -1)
126
+ speeds = out[..., 2] * sv
127
+ return mask, points, speeds