How the Go2 demo works

The demo runs a real MuJoCo simulation of a Unitree Go2 in your browser. A neural network policy I trained reads the robot's state 50 times a second and decides where its 12 joints should go. Nothing runs on a server.

Training

• Task. MuJoCo Playground ships a joystick locomotion task for the Go1, but not for the Go2. I ported it to the Go2 model from MuJoCo Menagerie and kept the reward, the 48-value observation and the action format unchanged.
• Pushes. Training includes random shoves to the trunk (kicks of up to 3 m/s), so the policy learns to recover instead of only walking on a perfect floor.
• Run. PPO with Brax on one RTX 4090: 200 million simulation steps in about 5.5 minutes.
• Gains. The Go2 is heavier than the Go1 (15.2 kg vs 12.7), so the joint stiffness went up from 35 to 50. At the Go1 value the standing pose sagged noticeably.

Checking it before it reaches the browser

The policy trains on MJX (MuJoCo on JAX/GPU) but runs on MuJoCo 4.0's CPU build in the browser. Small differences between the two can make a policy flail, so every export goes through these checks:

| Check | Result | |---|---| | Browser forward pass vs the training framework, 1,000 recorded steps | max difference 8×10⁻⁷ | | Walk 60 s at 1 m/s in CPU MuJoCo | 58.7 m, no fall | | Pushes of 0.5–3 m/s, standing and walking, 10 directions each | 120 of 120 recovered |

In the browser

• Physics. The official @mujoco/mujoco WebAssembly package, stepping every 4 ms. The policy runs on every fifth step, which is the timing it was trained with.
• Policy. No inference runtime. The network is small (48 inputs, hidden layers of 512, 256 and 128, 12 joint targets out), so a forty-line TypeScript forward pass replaces ONNX Runtime Web, whose WebAssembly alone is several MB. The weights are 0.7 MB.
• Rendering. Three.js draws the robot's meshes at the body poses MuJoCo computes each frame.
• Brain off zeroes the motor gains, so the robot collapses under the same physics. That's how you can tell the walking is learned control rather than an animation.

What broke along the way

• Loading took 8.6 seconds. MuJoCo was processing about 200k triangles of visual-only meshes. A build step now splits the model: physics gets collision shapes only (every body has explicit inertia, so the dynamics don't change), and the meshes go straight to Three.js. Compile time dropped to about 60 ms.
• The simulation ran 15% fast. Each frame took one physics step too many. A time accumulator that carries the leftover time between frames fixed it.
• The limp robot sank through the floor. The training model only lets the feet touch the ground. Trunk–floor contact is now part of the model, and it was on during training.
• The motion looked juddery. Running five physics steps at once on each policy tick made the picture update at 50 Hz. Physics now steps one 4 ms step at a time.

Credits

• MuJoCo and MuJoCo Playground (Apache-2.0) from Google DeepMind.
• The Unitree Go2 model from MuJoCo Menagerie (BSD-3-Clause).
• Prior art: unitree-mujoco-web runs Unitree policies in the browser on the same stack.