Training a humanoid on a MacBook

Whole-body control for humanoids is usually a data-center job. The standard recipe is model-free reinforcement learning over thousands of GPU-simulated robots, and the SuperMimic paper puts a number on it: its baseline trains for 7 days on 128 GPUs.

SuperMimic's own answer is a different algorithm. Learn a differentiable world model of the robot from simulator data, then improve the policy by backpropagating through short rollouts imagined inside that model. The paper reports beating the 128-GPU baseline in about 3 hours on a single GPU. XMimic, by Nico Bohlinger and Jan Peters, rebuilt that recipe in JAX and tracks the 40 LAFAN1 clips in under a minute on four H100s.

All of it assumes CUDA. The simulator is MuJoCo Warp, the learner leans on bf16 tensor cores, and whole training iterations are compiled into single GPU programs. A MacBook has none of that. What it does have is eight CPU cores, a GPU that shares memory with them, and MLX.

The method is sample-efficient, and that changes where the compute goes. Each training iteration needs only about 2,500 new simulated steps, because the policy learns mostly from imagination: 512 rollouts of 32 steps through the world model, differentiated end to end. So the physics can be slow and the learning has to be fast, which is exactly the shape of a laptop.

• CPU, physics. MuJoCo's plain C engine steps 1,024 robots in a pool of eight threads, about 22,000 control steps per second. Every robot keeps its own friction and center-of-mass randomization.
• GPU, learning. The world model, the policy, the critic, the rewards and the forward kinematics all run in MLX. The replay buffer is a numpy array, which costs nothing to hand over because host and GPU share memory.

A port that "seems to train" is worth little, so before the first real run I wrote a test that feeds the same inputs to the original JAX code and to the MLX port and compares the outputs: rotations, kinematics, motion features, rewards, networks, optimizers, and finally the gradients themselves.

All 351 checks pass. The gradient of the policy loss through 32 imagined steps agrees with JAX to a relative error of 2e-6, in float32 on both sides.

Metal hates small matrices

The first working version took 2.5 seconds per iteration, and the big matrix multiplies were not the problem. The 3x3 rotation products were. Thousands of them run in the kinematics and the reward, and both obvious ways to write one are slow on Apple's GPU: a batched matmul of tiny matrices, and a broadcast product summed over an inner axis.

# slow on Metal: a broadcast product summed over an inner axis
def mm3(A, B):
    return (A[..., :, :, None] * B[..., None, :, :]).sum(-2)

# 10-20x faster: three column-row products, fused into one kernel
def mm3(A, B):
    return (A[..., :, 0:1] * B[..., 0:1, :]
          + A[..., :, 1:2] * B[..., 1:2, :]
          + A[..., :, 2:3] * B[..., 2:3, :])

Automatic differentiation then puts the slow reductions back into the backward pass, so the products also got hand-written gradients, which are products of the same kind. Together with precomputing the reference motion's kinematics once, forward kinematics with gradients for 16,000 robot poses dropped to 18 ms and an iteration to about 1.3 seconds.

Two compiler bugs

mx.compile is what makes those formulations fast, and it failed twice. Splitting the same intermediate array twice crashes it with a missing-key error, so the rotation code uses slices instead. Then the fused gradient kernels grew past Metal's limit on kernel arguments, which a no-op mx.contiguous on each custom gradient fixes by acting as a fusion barrier.

The NaN at iteration 620

The first long run looked healthy for 16 minutes. Then every policy update started getting skipped as non-finite, and the policy silently stopped improving. The loss was finite. Only its gradient was NaN.

# mlx.nn.elu
mx.where(x > 0, x, mx.exp(x) - 1)        # gradient is NaN for x > 88

# what the trainer uses now
mx.where(x > 0, x, mx.expm1(mx.where(x > 0, 0.0, x)))

MLX's built-in ELU evaluates exp(x) on the branch it does not select. For a pre-activation above 88 that overflows in float32, and a zero cotangent times infinity is NaN. One imagined rollout with a large activation was enough to poison the critic for good. JAX's ELU guards against this; the port now does too.

16 GB is not much

MLX keeps freed GPU buffers in a cache whose default limit is nearly all of the machine's memory. The trainer grew to 7.5 GB, I started a second process next to it, and the laptop went into swap and stopped making progress. Capping the cache at 1 GB and giving each physics thread a copy of the robot with its meshes reduced to convex hulls (15 MB instead of 75 MB, same contact dynamics) brought the footprint to 2.6 GB at the same speed.

The run trained for two hours on the 40 G1-retargeted LAFAN1 clips, about 2.4 hours of motion. Each checkpoint is scored by playing every clip from its first frame without noise; a clip fails if the robot strays too far from the reference in height or tilt.

The target used for the H100 runs is a return of at least 0.82 with no failed clip, held for three checkpoints. The MacBook gets there at iteration 3,000, after 79 minutes. The hardest clips are the fall-and-get-up sequences: a first run that I stopped at 37 minutes failed all six of them, and the full run completes them.

The honest fine print:

• An H100 reaches the same target in about 2.5 minutes. The laptop is roughly 30 times slower, at a quarter of the batch sizes.
• The policy tracks body pose closely (4 to 7 degrees of mean joint error) but its position in the world drifts 0.1 to 0.7 m from the reference over 12 seconds.
• My checkpoints are 250 iterations apart and each clip is evaluated in one deterministic world, a coarser protocol than the published one.
• The SuperMimic numbers I quote are from the paper's abstract and are for a far larger motion benchmark than these 40 clips.
• This is simulation only. Nothing here has touched a real robot.

Not because a laptop should replace a cluster. It is because the loop of changing something and seeing a humanoid move differently now fits on the machine I already carry around, with no queue and no cloud bill. An algorithm that needs a thousand times less compute stops being a data-center experiment and becomes something you can iterate on from a couch.

I did the port in one long pairing session with Claude Code, which wrote most of the code while I steered, and the equivalence test is what let me trust the result.