WorkCell: five ways to teach one simulated robot arm

I built the simulated cell first: the MuJoCo scene with the arm, block, tray and two cameras, and an inverse-kinematics solver, the maths that converts "put the gripper at this point, pointing down" into the five arm joint angles that achieve it. Then I wrote the scripted controller that produces the demonstrations. It reads the block's true position from the simulator and moves the gripper through a fixed sequence of positions (above the block, down onto it, close the gripper, lift, over the tray, open, pull away), completing the task in 99 to 100% of recording attempts; failed recordings were discarded. It is the teacher: the model being trained never sees the block's position, only the two camera images and its own six joint readings. Each demonstration (the two camera images, the six joint readings and the motor command, 15 times a second) was saved in the dataset format of LeRobot, Hugging Face's open-source robot-learning library, so that LeRobot's training code could read it and so that it could be published.

Then I trained ACT, the imitation model introduced in the 2023 paper I compare against in section 05. It takes the two camera images and the six joint readings and outputs the next 45 motor commands as one plan (3 seconds of motion), of which the first 15 are executed before it looks again. I trained it four times, changing only the demonstrations: 100 recorded with the lighting, colours and physics always the same (fixed conditions); 160 recorded the same way except that 20 of them (12.5%) deliberately miss the block by about a centimetre on the first grab and then re-grab it; 150 recorded under randomly varied lighting, scene colours, block weight and friction (this is called domain randomization); and 210 randomized ones, about 30% of them with the block placed in the band of the table nearest the arm, where the 150-demonstration model kept failing. Every model was tested on the same list of block positions, on both the 20-attempt and the 50-attempt test.

The demonstrations, not the model, decided the outcome. The model trained on the 100 fixed-condition demonstrations succeeded in 37 of 50 test attempts (74%). The model trained on the 150 randomized demonstrations succeeded in 45 of 50 (90%) on the same fixed-condition test, so training under randomized conditions cost nothing on the unrandomized test; it added 16 percentage points. When I randomized the test conditions too (15 attempts per condition, so those figures are rough), that model scored 13 of 15 with nothing changed, 13 of 15 with the lighting and colours changed, 13 of 15 with the physics changed, and 15 of 15 with both changed. The paper's plan-averaging technique (temporal ensembling: make a new 45-command plan after every single command and average the overlapping plans; section 05) raised the fixed-condition model from 13 to 15 of 20 on the 20-attempt test. In short: one model design, trained four ways, scored anywhere from 62% to 90% depending only on what it was shown.

Training on the 160 demonstrations that include miss-and-re-grab, which I expected to teach the model to recover from a bad grab, gave 31 of 50 (62%) against the 100-demonstration model's 37 of 50 (74%); on the 20-attempt test the two scored 10 and 13, and switching plan averaging on took the re-grab model from 10 to 8. The reason: at the moment of a grab, the recordings now contain two different correct next moves, lift or go back and retry, and a model that averages what it has seen produces a movement that is neither. The 210-demonstration model fixed the failures near the arm: on a separate probe of 12 grid positions there it succeeded on 12 of 12 where the 150-demonstration model had managed 5 of 12. But on a second 12-position probe at the far side of the table it succeeded on only 4 of 12, failing all 8 positions furthest from the arm, where the 150-demonstration model had managed 12 of 12; and on the 50-attempt test it scored 44 of 50 (88%) against the earlier 45 of 50. The retraining moved the failures instead of removing them. Both of those models (the re-grab one and the 210 one) are published on Hugging Face, labelled as negative results, with a description page saying what went wrong.

Instead of training a model from scratch, I started from SmolVLA, a freely downloadable model with 450 million parameters (the adjustable numbers inside a model that training sets; more parameters means a bigger model), published by Hugging Face. Hugging Face pretrained it on a large collection of community recordings, mostly from real SO-100-family arms, the same family as mine. It takes camera images, joint readings and a text prompt (here always "Pick up the red block and place it in the blue tray") and outputs motor commands. A model this size just fits on Kaggle's free 16 GB GPU for fine-tuning. Fine-tuning here meant continuing its training on my demonstrations for 12,000 training steps (each step shows the model one batch of recorded moments and adjusts it slightly), about 8 hours on a free Kaggle GPU (the 20,000 I planned would have exceeded Kaggle's 12-hour session limit), while leaving the parts of the model that read the images and the text unchanged (frozen) and adjusting only the part that turns them into motor commands. I tested it exactly the way I tested ACT: the same 20 block positions, the same 16-second limit, the same definition of success.

The unmodified pretrained model succeeded in 0 of 20 attempts: a model trained on real-robot recordings does not work in a simulated cell with cameras in different places until it is fine-tuned. After fine-tuning on the 100 fixed-condition demonstrations it succeeded in 11 of 20 (55%). After fine-tuning, with identical settings, on the 160 demonstrations that include the miss-and-re-grab ones, it succeeded in 18 of 20 (90%), matching the best ACT model. Those same 160 demonstrations had lowered ACT from 13 to 10 of 20 on this 20-attempt test (37 to 31 on the 50-attempt one), so identical data moved the two models in opposite directions.

My explanation is that the part of SmolVLA that generates the motor commands is trained with a method called flow matching, which produces one movement drawn from a range of possibilities, whereas ACT is trained to minimise its average error against the recordings, which pulls it toward the midpoint when two recordings disagree; so SmolVLA can hold "lift" and "re-grab" as two separate options where ACT blends them. That fits the result but is not proven: SmolVLA has 450 million pretrained parameters and ACT has 52 million trained from scratch, so model size and pretraining are alternative explanations I did not rule out. The model accepts a text prompt, but every demonstration used the same sentence, so I have not tested whether it responds to the words at all. Each result is one training run and 20 test attempts, copied by hand from the program's on-screen output rather than from a saved log file. The code in this repository is the fine-tuning notebook and the test script; the simulator, the scripted controller, the demonstrations and the test are reused from Part A.

A depth camera above the table records, for every pixel, both its colour and its distance from the camera. Knowing the distance to each pixel lets the program convert the image into a point cloud: a list of 3D points on every surface the camera can see. The scene has three blocks of the same size and shape (24 mm cubes): a red target and green and yellow look-alikes. The program works through four steps, each an algorithm I implemented myself. The conversion from depth image to point cloud is the first. RANSAC then finds the table: it fits a flat plane to small random samples of points, keeps the plane that the most points agree with, and removes those points. DBSCAN splits the remaining points into separate objects: points packed closely together count as one object and a gap starts a new one, so it never needs to be told how many objects there are; the object whose points are red is the target. ICP then fits an ideal 24 mm cube (a set of points on the surface of a perfect cube, not a trained model) to those points: it repeatedly pairs each measured point with the nearest point on the ideal cube and shifts and rotates the cube to fit better, until it stops improving. The result is the block's position and its rotation about the vertical axis (its yaw), which is all the gripper needs, because a cube resting flat on the table can only differ in how far it is turned. The program is also not told where the camera is relative to the arm. I measured that (a camera calibration) by placing a small sphere at 11 known positions in the simulator and solving for the camera position and orientation that best explains where the sphere appeared in each image. All of it is about 1,030 lines of Python using NumPy, the standard array-maths library; the only ready-made routine is SciPy's nearest-neighbour lookup (finding the closest of thousands of points quickly).

In the full test (depth image in, block in tray out) the arm picked the red block into the tray in 20 of 20 layouts in which all three blocks were placed at random. The block's position was located to within about half a millimetre: 0.58 mm on average over 15 layouts in Part C's own test, and 0.45 mm on average (worst case 0.77 mm) over the 12 saved sets of block points in the benchmark file in the Part E repository, which runs this exact code. For comparison, simply taking the centre of the block's visible points is off by about 8 mm in my test, because the camera only sees the top and the front faces. The rotation was found to within about 6 degrees on average, and the gripper still closes on the block correctly even when the rotation estimate is up to 20 degrees wrong. The camera calibration recovered the camera's true position to within 0.32 mm and its orientation to within 0.011 degrees; the calibration file that shows this is in the repository. Nothing is trained and no GPU is used. The final step, the rotation search and ICP fit, takes about 14 ms per block on a laptop CPU; the table-finding and clustering steps before it are not timed.

ICP needs a starting guess for the rotation. My first way of choosing it used PCA, a standard method for finding an object's longest axis. A cube seen from above has no longest axis (a square is equally wide in both directions), so the "axis" it found was random noise and ICP kept locking onto a fit 45 degrees off. I replaced it with a brute-force check: try 18 candidate rotations, every 5 degrees across the 90 degrees a cube can differ by, keep the one that fits best, and start ICP from there. Two formulas I derived by hand for how rotating the wrist turns the gripper's jaws each had the sign wrong, so the program now measures which way the jaws point before every grasp, by setting the wrist angle in the arm's model and reading where the jaws end up, instead of trusting a formula. And lowering the gripper straight down onto the block sometimes landed the gripper's fixed jaw (one jaw moves, the other does not) on top of it, because there is only a couple of millimetres of clearance between that jaw and the block, so the gripper now comes in from the side. The limitation of the whole approach: it only works because the block's shape and size are known and the simulated depth camera has no noise. A new object means new code, and a real camera would need noise handling I have not written.

Reinforcement learning works like this: the arm starts by moving at random in the simulator, after every move it receives a score (the reward), and a learning algorithm adjusts the arm's controller (for PPO and SAC, a neural network that picks the next motor command from the joint readings and the position of the target or block) so that the kinds of moves that earned higher scores become more likely. Repeated over many thousands of moves, the controller improves. I set two tasks, both using the simulator's exact numbers (joint readings and block position) as the input rather than camera images. Reach: move the gripper tip to within 2 cm of a target point that changes every attempt (each target is first checked to be reachable). The reward for reach was minus the distance in metres from the gripper tip to the target, so the score rises as the tip gets closer, minus a small penalty for large motor commands, plus 1 for arriving within the 2 cm. Lift: grasp the block and hold it more than 7 cm above the table for the rest of the attempt; its final reward, after four versions (below), pays steeply for the first 7 cm of lift, then 5 points for every move the block is held above 7 cm, minus a penalty above 14 cm, plus small amounts for getting the gripper near the block and closing on it.

On reach I compared three methods. Two are the standard reinforcement-learning algorithms PPO and SAC, taken from Stable-Baselines3, the well-tested library that implements them (PPO learns only from its most recent moves; SAC stores every move it has ever made and keeps learning from all of them, which usually means it needs fewer moves to learn). The third I wrote myself in about 150 lines of NumPy and PyTorch (the standard neural-network library). First, three small neural networks (their predictions averaged) learn to predict where the arm's joints will be after any given command, from 8,000 random moves. Then, before each real move, a planner generates 256 candidate sequences of 15 moves, scores each sequence with the same reward, using the networks' predictions of where the tip will end up instead of running the simulator, keeps the best 32, generates new candidates around those, repeats four times (this search is the cross-entropy method), and executes only the first move of the winning sequence before planning again. A command simply nudges each servo's target position by a known amount, so I compute that bookkeeping exactly and let the networks learn only how the arm actually moves in response. On lift I used SAC alone.

My planner (one training run, against three each for SAC and PPO) reached the target in 20 of 20 test attempts after learning from just those 8,000 random moves, and stayed at 20 of 20 after three further rounds, each round being a few thousand more moves made by the planner itself, added to the training data, with the networks retrained (the plot of success per round is in the repository; the raw evaluation files are not). SAC reached 19 or 20 of 20 on each of three training runs at its 100,000-move budget, and PPO reached 20 of 20 on all three runs at its 300,000-move budget; the reward-over-time curves in the repository show SAC levelling off around 50,000 moves and PPO around 250,000, with PPO still well below its final level at 100,000. On lift, the final reward got SAC to 18, 15 and 11 successes out of 20 (90%, 75%, 55%) across three separate training runs, holding the block around 16 cm up rather than as high as the arm can reach.

My first lift results were invalid because of a bug in the code that resets the simulator between attempts: every attempt after the first ended after one move. I caught it because the average attempt length came out at exactly 8.45 moves; attempts are capped at 150 moves, so (150 + 19) / 20 is one full 150-move attempt and nineteen one-move attempts. Then the reward had to be rewritten three times. Version one gave points every move for keeping the gripper close to the block and holding it, plus a one-time bonus for lifting the block past 7 cm, after which the attempt ended; the arm learned to hover the block just below 7 cm for the whole attempt, because collecting the per-move points for 150 moves was worth more than the bonus. Version two gave points for every move with the block above 7 cm, and nothing discouraged going higher, so the arm yanked the block as high as it could reach. Version three made the points rise gradually from table level up to 14 cm, which made the reward for the first few centimetres too small to be worth grasping at all, so the arm often never lifted the block. The final version is the one described above. (The three earlier versions are described from my development notes; their code has been replaced.) Each SAC lift run also spent a long stretch grasping the block without lifting it before suddenly improving, at a different point in each run, so one of the three lift runs needed 500,000 moves to get there where the other two needed 300,000. And the prediction networks in my planner cannot tell when one of their predictions is unreliable, and were used only for reach, never for lift.

Parts A to D all ran inside one physics simulator. Real robots run on ROS 2: each part of the robot's software (the camera reader, the motor driver, the route planner, and so on) is a separate program, and ROS 2 is the messaging system they use to send data to each other. It is the industry standard, and ready-made mapping and navigation programs exist for it. I set it up under WSL2 on my laptop. I rewrote the final step of Part C's block-locating code, the rotation search and the ICP fit, in C++ using Eigen, the standard C++ matrix library; the table-finding and clustering steps stayed in Python. I also wrote my own KD-tree (the data structure that makes finding the nearest point among thousands fast) in 120 lines, rather than pulling in a large point-cloud library. The rewrite runs as a ROS 2 service: another program sends it the block's 3D points and gets back the block's position and rotation. To check it, I saved 12 sets of block points produced by Part C's depth camera and clustering steps, along with the true block positions from the simulator, ran both the Python original and the C++ rewrite on all 12, and compared the outputs to the last decimal place.

A production line also has to move parts between cells, so for that I used a simulated TurtleBot3, a small wheeled robot that ships with an official simulation model. A small exploration program I wrote drove it around a simulated room it had no map of for 240 seconds (keep going while the way ahead is clear, turn toward the more open side, with an occasional random turn), while slam_toolbox, a ROS 2 mapping program, built a map from the robot's lidar (a spinning laser that measures the distance to the walls around it). Then Nav2, the ROS 2 navigation program, using AMCL (a program that matches each new lidar scan against the map to work out where the robot is), drove it to six target points: five picked by a script I wrote from the open areas of the map, plus a return to the start.

For all 12 sets of block points the C++ rewrite produced the same block position as the Python original to within 0.001 mm (the smallest difference the benchmark records: 10 of the 12 were identical to that digit and 2 differed by 0.001 mm), the same rotation, and the same number of ICP iterations. Both versions put the block's centre within 0.45 mm of its true position on average. The C++ version took 4.33 ms per set against Python's 14.0 ms, 3.3 times faster on average (the average of the per-set speed-ups, which ranged from 2.4 to 3.9), on a single CPU core. The navigation run reached all six target points, taking 13 to 49 seconds per point. The benchmark results, the timings of my first slower C++ attempt, and the navigation log are all files in the repository. The demo animations were drawn afterwards from the robot's logged positions and the map, rather than screen-captured, because no screen recorder works with WSL2's graphics.

My first working C++ version was only 1.5 to 2 times faster than Python (9.09 ms per set; that timing is also in the repository), because the Python version already uses SciPy's nearest-neighbour search, which is itself written in C. The real speed-up came from a change of approach: build the KD-tree for the ideal cube once and, in each ICP iteration, apply the current shift and rotation to the measured points (so they are expressed in the cube's own coordinates) instead of rebuilding the KD-tree around the moved cube every iteration. Under WSL2, most failures are silent: nothing crashes, no error appears, the robot simply does not move. I wrote up nine of them with dates, each with the symptom, the cause and the fix: for example, a speed command sent in an older message format that the robot's simulation bridge silently ignores, a sender and receiver whose delivery settings did not match so no data ever arrived, and a simulator that froze silently while the mapping program kept re-sending its last map, so every process looked alive. What the benchmark does not show: it checks that the two versions agree on the rotation, but it reports accuracy against the true block position only, not against the true rotation (about 6 degrees in Part C's own test). The KD-tree has no separate correctness tests beyond agreeing with SciPy on those 12 sets; the 4.33 ms is the calculation alone, timed by a standalone benchmark program, and does not include the time to send a request through ROS 2 and get the answer back; and a second navigation run that also reached 6 of 6 was not saved, so this page counts only the recorded one.