MuJoCo and CleanRL: train a balance bot with PPO

Learn to build a Gymnasium wrapper and train a balance bot with PPO on CPU, using two training phases and an actor network with two 16-node hidden layers.

Player not loading? Watch on YouTube

This tutorial builds on an earlier MuJoCo robot simulation to train a balancing policy with Python. The speaker creates a Farama Gymnasium environment and adapts CleanRL's continuous-action PPO implementation. All training and evaluation in this episode happen in simulation; deployment to the ESP32-based robot comes later.

The environment supplies four observations: pitch, pitch rate and the two wheel velocities. The agent returns two normalized wheel torque commands. Simulation advances in 5-millisecond steps, with a default tipping threshold of 30 degrees. The reward combines an upright bonus with penalties for pitch and motor effort. Position and yaw penalties use simulator information that the speaker says would not be available on the real robot.

Training follows a two-phase curriculum. First, the bot learns to balance while it can wander. The second phase adds penalties for leaving the origin and turning. The example uses four parallel environments on CPU, with two 16-node hidden layers in each actor and critic network.

The walkthrough covers environment checks, checkpoints and TensorBoard plots. The speaker flags a suspected logging bug in explained variance and uses value loss to assess the critic instead. It ends with an actor-only ONNX export and inspection in Netron. The speaker reports successful simulated balancing but warns that this initial policy performs worse on real hardware.