In plain words
The problem in one picture: a rendered cube, and a network that has to say where it sits and how it is turned.
A robot that is going to pick something up needs to know exactly where that thing is and how it is rotated. This project builds the smallest honest version of that problem: a single red-ish cube on a flat table, one camera looking straight down, and the question "give me x, y and the rotation angle".
There are two obvious ways to answer it.
- Write the rule yourself. The cube is a coloured blob on a plain table. Threshold the image on colour, find the blob, take its centre, fit the smallest rectangle around it and read off the angle. Thirty lines of OpenCV, no training, runs in 0.05 ms.
- Train a network. Render thousands of labelled images, feed them to a small convolutional network, and let it learn the mapping.
Almost every tutorial does option 2 and declares victory. This project does option 1 first, and then does it a second time, better, before the network is allowed to compete. That ordering is the whole experiment, and it changes the answer: the first hand-written baseline loses to the network by 4.5×, and after one calibration constant is fitted to it, more than half of that gap disappears.
Building the strong baseline after the model would have inflated the network's apparent advantage by roughly 2×. What the network actually wins is not the typical case — it is the tail: the classical method fails rarely and badly, the network fails never and mildly.
Why simulation
Labelling real photographs by hand is slow, and a human-placed label of "the cube centre is at x = 4.2 cm" is not accurate to a millimetre. In MuJoCo the numbers go the other way: the position and angle are typed in, and then the image is rendered from them. The label is not an annotation, it is the input. Twelve thousand perfectly labelled images take about twelve seconds to generate.
Terms used on this page
The task, defined precisely
What pose means here, what the network outputs, and the symmetry that makes yaw awkward.
Position
The cube centre, in metres, in the table frame, sampled uniformly in ±0.075 m. The camera is directly overhead and its image plane is parallel to the table, which means the mapping from pixels to metres is exactly affine — no perspective correction, no distortion model. At 128 px the scale is 367.9 px/m, i.e. one pixel is 2.72 mm.
Rotation, and the trap in it
A cube seen from directly above shows a square. A square looks identical at 5° and at 95°. So the rotation is only observable modulo 90°, and asking a network for "the angle in degrees" is asking it for something that does not exist.
Worse, regressing a raw angle puts a discontinuity in the middle of the label range: 89° and 1° are nearly the same pose but numerically far apart, so the network hedges toward the middle and is systematically wrong near the boundary. The fix used here — and it is standard practice in pose estimation — is to predict a point on a circle instead of an angle:
target = (sin 4θ, cos 4θ) # YAW_FOLD = 4Multiplying by four folds a full turn onto four identical quarters, so θ and θ+90° map to exactly the
same target and there is no seam anywhere. Decoding runs it backwards:
θ = atan2(s, c) / 4, wrapped into [0°, 90°).
Crucially, OpenCV's minAreaRect has exactly the same ambiguity. Both methods are scored
under the same 90° fold, by the same function, so the comparison is fair by construction rather than
by assertion.
Two appearance regimes
| Regime | What varies | What it tests |
|---|---|---|
easy | Nothing. Fixed cube colour, fixed light position and intensity, fixed table shade. | Whether a hand-written rule is simply the better engineering choice when the world is controlled. |
hard | Cube hue uniform over the whole colour wheel, saturation and value randomised; table shade randomised; light position and intensity randomised. | Whether a fixed prior survives appearance change — the miniature version of the sim-to-real question. |
Note what is not randomised: geometry, camera pose, object size. Only appearance moves. That keeps the comparison to a single axis, so a result can be attributed.
The pipeline
The order the programs run in, from scene file to exported model.
Steps 2 and 4 are the same classical method. It appears twice on purpose — once before the network so the network has a target, and once after so the network's win can be attributed. Step 4 is the one that changed the conclusion.
Every file, and what it is for
Every source file, how big it is, and the single job it owns.
| File | Lines | What it does |
|---|---|---|
| common.py | 116 | The single source of truth. Scene constants that must match the MJCF; the yaw
fold and its inverse; the pixel↔world projection; the dataset loader; and — most importantly —
pose_metrics(), the one scoring function used by every method. If the
classical baseline and the network were scored by two implementations they would drift apart, and
the comparison the whole project exists to make would be worthless. |
| scene.xml | 24 | The MuJoCo scene. A plane at z = 0, one light, one overhead camera at 0.45 m with a 45° vertical
field of view and axes aligned to the world, and a box of half-size 0.03 m on a free joint.
Shadows are switched off (shadowsize="0") because shadow mapping is the expensive
part of software rendering. |
| gen_dataset.py | 132 | Renders the dataset. Samples x, y and yaw, writes them into qpos, mutates the
cube/table colours and the light according to the regime, calls mj_forward, renders,
and stores the image into a memory-mapped .npy array. Writes a
meta.json recording image size, px/m, and the label column order. Train and
validation splits use different seed streams and are never mixed. |
| view_dataset.py | 52 | Label verification. Projects each ground-truth label back into pixel space through the analytic overhead projection and draws a marker. If the markers do not land on the cubes, the label convention is wrong — and every number produced afterwards would be consistently, silently wrong. |
| baseline_cv.py | 101 | Classical baseline, two variants. --method red is a fixed hue window — the naive
prior "the cube is red". --method sat keeps anything sufficiently saturated — the
weaker, more honest prior "the cube is the most colourful thing on a near-grey table". Both then
run: morphological open → largest contour → area floor → moments centroid +
minAreaRect angle → convert to world units. |
| baseline_v2.py | 60 | The fair baseline. Runs the same detector over the train split, fits one scalar radial correction by least squares, applies it unchanged to validation. One learned parameter against the network's 27 000. |
| model.py | 86 | Two heads on one convolutional trunk: a generic flatten head and a spatial soft-argmax keypoint head. Described in detail below. |
| train.py | 145 | The training loop. Normalises the four targets onto a common scale, uses SmoothL1 loss, AdamW at 2e-3 with a cosine schedule, batch 64, 40 epochs, no augmentation. Prints per-epoch throughput and flags a sustained >20 % drop, because this machine exposes no CPU temperature sensor. Reports the final epoch and selects on nothing. |
| export_onnx.py | 100 | Exports to ONNX and then proves the export three ways: max absolute tensor difference on a real batch, the full task metrics recomputed through ONNX Runtime, and batch-1 latency at both 1 and 8 threads. |
| test_common.py | 67 | Tier-1 tests. The yaw fold, the projection and the metric checked against values derived on paper. Needs no dataset and no weights — it runs from a fresh clone in seconds. |
| verify.sh | 55 | Reproduces the README's claims from a clean checkout, cheapest tier first. Pins
MUJOCO_GL and the thread count because both change the numbers, and skips the data
tiers with instructions rather than failing when data/ is absent. |
| requirements.txt | 9 | Explicit CPU wheels: torch==2.14.0+cpu. Pinned because every number in the README
was measured against them. |
| runs/ · out/ · data/ | — | Three run directories with tracked final.pt checkpoints and
metrics.json; qualitative PNG sheets; and the generated dataset, which is gitignored
at ~1.3 GB and regenerates in about 30 seconds. |
Making the data
How the training images are generated in MuJoCo, and what is varied between them.
Where the label comes from
The generator does not run physics. It writes the pose directly into the simulator's state vector and
calls mj_forward, which recomputes all the derived quantities without advancing time:
data.qpos[:3] = [x, y, CUBE_HALF]
data.qpos[3:7] = [cos(yaw/2), 0, 0, sin(yaw/2)] # quaternion about z
mujoco.mj_forward(model, data)
renderer.update_scene(data, camera="top")
imgs[i] = renderer.render()
labels[i] = [x, y, sin(4·yaw), cos(4·yaw), yaw]A quaternion about the z axis is (cos(θ/2), 0, 0, sin(θ/2)) — the halved angle is the
standard quaternion convention and is the second most common place to introduce a silent factor-of-two
error in this kind of code. The label array keeps the raw yaw in a fifth column for debugging only;
nothing trains on it.
Randomising appearance
On hard, four things move per image:
These are written directly into model.geom_rgba, model.light_pos and
model.light_diffuse — fields of the compiled model, mutable at run time, so one XML covers
every appearance.
Memory-mapped storage
Twelve thousand 128×128 RGB images is about 590 MB. The generator opens the array with
np.lib.format.open_memmap and writes into it row by row, so the whole set never has to fit
in RAM at once. Training loads it the same way and sorts each batch's indices before slicing, which
turns random access into sequential reads from the file — a small change that matters when the data is
on disk rather than in memory.
Cost, measured
900–1080 images per second at 128 px, so 12 000 images takes about twelve seconds. Rendering is not the bottleneck in this project, which is worth stating because it usually is: on this machine OpenGL is software-rendered (llvmpipe under WSL) and shadows are disabled precisely to keep it that way.
The baselines — and why there are two of them
The classical-vision baseline, the bias found in it, and the one-line correction that removed it.
Version 1: two hand-written priors
Both variants share the same pipeline and differ only in the mask:
hsv = cvtColor(img, RGB2HSV)
m = MASKS[method](hsv) # 'red' or 'sat'
m = morphologyEx(m, MORPH_OPEN, ones(3,3)) # kill speckle
c = max(findContours(m), key=contourArea) # largest blob
if contourArea(c) < 30: return None # detection failed
(u,v),(w,h),ang = minAreaRect(c)
x, y = pixel_to_world(u, v, size)
yaw = mod(radians(-ang), pi/2) # image v grows down → sign flipThe two priors behave very differently once appearance moves:
| Prior | Detection rate on hard | What happened |
|---|---|---|
red — fixed hue window | 10.9 % | Missed 89 % of the set, because the cube is no longer red. Its accuracy on the 11 % it did find is unchanged, so its own error metric looks perfectly healthy. |
sat — most saturated object | 100 % | Fully robust to the appearance shift, with the same error as on easy. No network involved. |
A wrong prior does not raise an error. It returns fewer answers, and the answers it does return are fine. Detection rate and accuracy have to be reported together or the failure is invisible — which is exactly what happens when a paper reports only mean error.
Version 2: one scalar, fitted honestly
Version 1 has a systematic +2.98 mm radial bias — every prediction sits slightly too far from the table centre. The cause is geometric, not statistical: an overhead camera does not only see the cube's top face, it also sees the outward-facing side walls. The silhouette is therefore wider on the outward side, and its centroid sits outside the true projected centre.
That is not a fair fight. The network saw 12 000 labelled images; the baseline saw none. So the baseline is given the same privilege and no more: fit exactly one number, a radial scale correction, on the train split, and apply it unchanged to validation.
k = argmin ‖ k·r_measured − r_true ‖ # closed form, one line
= (r_meas · r_true).sum() / (r_meas²).sum()
= 0.951 # a −4.9 % radial correctionThe bias goes from +2.98 mm to +0.01 mm, and the median error falls from 3.41 mm to 1.91 mm — 53 % of the whole gap to the network, closed by a single scalar and no training loop.
The network — two heads, one question
Two output heads compared on identical data — flatten against spatial soft-argmax.
Why the usual head cannot be used
Nearly every image classifier ends with global average pooling: average the feature map over space, then classify. That works because "is this a cat" does not depend on where the cat is — translation invariance is a feature.
A pose regressor needs the opposite. The answer is the location. Average over space and you have thrown away the thing you were asked for. So the head has to preserve spatial information, and there are two ways to do that.
| Head | Trunk | Mechanism | Params |
|---|---|---|---|
flat | 4 stride-2 blocks, 128×128 → 8×8×128 | Flatten the whole feature map to 8192 values and put a single linear layer on it. No prior at all — the network has to discover the geometry of a raster from scratch. | 130 k |
softargmax | 3 stride-2 blocks, 128×128 → 16×16×64 | A 1×1 convolution produces 16 keypoint channels; each is turned into a probability map by a spatial softmax; the expected (u, v) of that map is computed. The layer's output is a coordinate. | 27 k |
Both then produce four numbers: normalised x, normalised y, sin 4θ, cos 4θ.
How spatial soft-argmax works
Given one feature channel of shape (H, W):
- Divide by a learned temperature and take a softmax over all H×W positions. The channel is now a probability distribution over pixels — "where in this image is my feature".
- Build two coordinate grids, u from −1 to 1 left-to-right and v from −1 to 1 top-to-bottom.
- Take the expectation:
E[u] = Σ p·u,E[v] = Σ p·v.
The result is a differentiable, sub-pixel coordinate. The temperature is a learnable parameter stored as its logarithm, so the exponential keeps it strictly positive without a constraint — a low temperature gives a sharp, argmax-like peak; a high one averages broadly.
This is the keypoint front end used by visuomotor policies — the deep-spatial-autoencoder line of work — and the reason to implement it here rather than read about it is that the project measures what the prior is worth: 5× fewer parameters, 2.6× faster, and more accurate, with a better tail.
Training details that are load-bearing
XY_RANGE so it lands in [−1, 1], matching the sine and cosine
which already do. Without this, 0.075 m of position error and 1.0 of sine error would be weighted a
thousand to one and the model would learn nothing but yaw.The first version of train.py kept the lowest-validation epoch and reported it. Two of the
three runs had already settled on the final epoch, so their numbers were unchanged; the third was
retrained and moved from 0.47 mm to 0.48 mm. The leak cost nothing measurable — which is the point of
checking rather than assuming.
Export, and what counts as proving it
ONNX export, the numerical tolerance it is checked against, and what a passing check does not prove.
ONNX is a portable graph format: train in PyTorch, run anywhere — including a C++ node on an embedded board, with no Python and no PyTorch installed. The file appearing on disk is not the finish line, though, because a graph can be numerically close and still have a permuted output order, a wrong normalisation, or a silently dropped layer.
So export_onnx.py checks three things:
- Tensor agreement. A real 256-image batch through both engines. Max absolute
difference:
6.6e-7. - Task-level agreement. The full validation split decoded exactly as training decodes
it, scored by the same
pose_metrics. Every figure identical to two decimals. - Latency at batch 1, at both 1 and 8 threads. Batch 1 is the only batch size a real perception node ever sees, and a ROS 2 node does not get the whole CPU.
On torch 2.14 torch.onnx.export defaults to the dynamo exporter, which imports
onnxscript and fails if it is absent. Passing dynamo=False selects the
TorchScript exporter and needs no extra dependency — worth knowing before adding a package to a
locked-down machine for no reason.
Results
The measured numbers: position error, yaw error, and runtime.
Validation split, n = 2000, independent seed stream, every row scored by the same function. "p95" is the 95th percentile — the tail, not the typical case.
| Method | Params | Regime | Detect | Median xy | p95 xy | Median yaw | p95 yaw | Radial bias | Latency |
|---|---|---|---|---|---|---|---|---|---|
| OpenCV, fixed red hue | 0 | hard | 10.9 % | 3.48 mm | 6.07 | 0.19° | 1.24 | +2.96 mm | 0.06 ms |
| OpenCV, saturation | 0 | easy | 100 % | 3.54 mm | 5.54 | 0.18° | 1.50 | +3.01 mm | 0.05 ms |
| OpenCV, saturation | 0 | hard | 100 % | 3.41 mm | 5.57 | 0.19° | 1.66 | +2.98 mm | 0.05 ms |
| v2 · OpenCV + 1 calibrated scalar | 1 | easy | 100 % | 1.94 mm | 2.84 | 0.18° | 1.50 | +0.16 mm | 0.04 ms |
| v2 · OpenCV + 1 calibrated scalar | 1 | hard | 100 % | 1.91 mm | 2.89 | 0.19° | 1.66 | +0.01 mm | 0.06 ms |
| CNN, flatten head | 130 k | hard | 100 % | 0.75 mm | 2.00 | 0.28° | 0.94 | +0.05 mm | 1.70 ms |
| CNN, soft-argmax head | 27 k | easy | 100 % | 0.48 mm | 1.04 | 0.16° | 0.51 | −0.01 mm | 0.65 ms |
| CNN, soft-argmax head | 27 k | hard | 100 % | 0.59 mm | 1.32 | 0.21° | 0.66 | −0.02 mm | 0.65 ms |
| same weights, ONNX Runtime, 1 thread | 27 k | hard | 100 % | 0.59 mm | 1.32 | 0.21° | 0.66 | −0.02 mm | 0.23 ms |
| same weights, ONNX Runtime, 8 threads | 27 k | hard | 100 % | 0.59 mm | 1.32 | 0.21° | 0.66 | −0.02 mm | 0.12 ms |
One pixel is 2.72 mm, and an early draft of the README wrote that down as a hard accuracy floor. It is not. It bounds a single-pixel measurement — but the cube covers roughly 480 silhouette pixels, and averaging over them shrinks the quantisation term by about √N, giving an aggregate floor near 0.12 mm. Both the calibrated baseline at 1.91 mm and the network at 0.59 mm legitimately sit below one pixel. Sub-pixel accuracy from a many-pixel object is expected, not a bug; reasoning about a floor from the wrong unit of measurement is how a correct result gets discarded as an error.
What the numbers actually say
Reading those results honestly — where the gain came from, and where it did not.
A wrong prior fails silently, and its own metric hides it
"The cube is red" misses 89 % of hard while its accuracy on the surviving 11 % looks
unchanged. Detection rate and accuracy must always be reported together, or the failure is invisible.
Robustness to appearance randomisation was free — and classical
Swapping "red" for "most saturated thing on a near-grey table" gives 100 % detection on hard
with the same error as easy. Domain-randomisation robustness is therefore not what
justified the network here. Picking a better hand-written feature was cheaper, and it worked.
Half the network's apparent win was calibration, not perception
Fitting one scalar on the train split takes the baseline from 3.41 mm to 1.91 mm — 53 % of the gap, closed by a single number. Building the strong baseline second would have inflated the network's result by roughly 2×.
What the network genuinely wins is the tail
After calibration the baseline's median yaw error (0.19°) is slightly better than the network's
(0.21°) — minAreaRect is near-exact on a clean mask. But its p95 is 1.66° against 0.66°, and
p95 position is 2.89 mm against 1.32 mm. The classical method fails rarely and badly; the network fails
never and mildly. For anything feeding a controller, the tail is the number that matters.
The architectural prior beat the parameter count
Flatten head: 130 k parameters, 1.70 ms, 0.75 mm. Soft-argmax head: 27 k parameters, 0.65 ms, 0.59 mm with a better tail. 5× smaller, 2.6× faster, and more accurate. Handing the network a coordinate beats making it learn what a coordinate is from a flattened grid.
The classical method is still 11× faster
0.06 ms against 0.65 ms per image. If 1.9 mm is inside tolerance, the network is the wrong engineering choice regardless of being three times more accurate.
The runtime mattered more than the model
Identical weights, identical outputs: PyTorch eager needs 0.65 ms on eight threads; ONNX Runtime needs 0.23 ms on one and 0.12 ms on eight. That is ~2.8× faster on one eighth of the cores, with nothing about the network changed. Before optimising an architecture for edge latency, check whether the framework is the cost. Here it was most of it — and the classical speed advantage shrinks from 11× to 3.8× as a result. A 116 KB ONNX file at 0.23 ms is a real deployment target for a C++ node.
Running it
Setup, the commands in order, and the runtimes to expect on a CPU.
git clone https://github.com/AungKaung1928/mujoco-cube-pose-cnn.git
cd mujoco-cube-pose-cnn
python3 -m venv .venv && . .venv/bin/activate
pip install -r requirements.txt
python gen_dataset.py --regime easy --n 12000 # ~12 s
python gen_dataset.py --regime hard --n 12000
python view_dataset.py --regime hard # check the labels first
python baseline_cv.py --regime hard --method red # the naive prior
python baseline_cv.py --regime hard --method sat # the honest prior
python baseline_v2.py --regime hard --method sat # the calibrated one
python train.py --regime hard --head softargmax --epochs 40
python export_onnx.py --regime hard --head softargmaxOr check every claim without retraining
./verify.shThree tiers, cheapest first. Tier 1 is test_common.py — the yaw fold, the projection and
the metric against paper-derived values, no data needed. Tiers 2 and 3 need data/
(gitignored, ~30 s to regenerate) and re-score the tracked checkpoints through both PyTorch and ONNX
Runtime. The script defaults MUJOCO_GL to glfw and pins the thread count,
because both change the numbers.
Measured environment
Intel Core Ultra 5 225H, 14 cores, no GPU, WSL2, software OpenGL (llvmpipe). Rendering 900–1080 img/s at 128 px; training ~1750 img/s on 8 threads, decaying 25–35 % in the second half of a run as the CPU hits its sustained power limit.
The short version
The whole project compressed into one paragraph.
Estimate the planar pose of a cube from one 128×128 overhead render, and measure when a learned model
beats hand-written geometry and when it does not. Twelve thousand images per regime with labels read
straight out of MuJoCo. Two hand-written OpenCV baselines, then a third with one scalar calibrated on
train — which closes 53 % of the gap and shows the first baseline had been unfairly weak. A 27 k-parameter
soft-argmax keypoint head reaches 0.59 mm median and 1.32 mm p95 on randomised appearance, beating a
130 k-parameter flatten head on every axis. Exported to ONNX and verified twice — tensors and task
metrics — where the same weights run 2.8× faster on one thread than PyTorch eager on eight. Every claim
re-runnable from a clean clone through a tiered verify.sh, and the two corrections made along
the way are left in the record rather than edited out.