Reinforcement Learning for Robotics Part 2: Train a Balance Bot with PPO

169

2026-08-06 | By ShawnHymel

Displays LCD / TFT Microcontrollers Robot Kits M5Stack

In the previous episode, we built a digital twin of the M5Stack BALA 2 Fire robot and imported it into the MuJoCo simulator. Now it's time to actually teach it something. In this episode, we develop a Gymnasium environment wrapper for our MuJoCo simulation, implement a custom version of Proximal Policy Optimization (PPO), and use a technique called curriculum learning to train the robot to balance on its own.

Everything in this episode happens in simulation. We won't deploy to the real robot until the next episode, which is also where we'll tackle the sim-to-real gap. For now, the goal is a trained agent that balances reliably in MuJoCo.

All project files are available here: github.com/ShawnHymel/reinforcement-learning-for-robotics.

RL Basics: A Quick Refresher

If you're new to reinforcement learning, I'd highly recommend watching this intro to RL video before continuing.

Image of Reinforcement Learning for Robotics Part 2: Train a Balance Bot with PPO

As a quick refresher: the robot periodically produces an observation vector describing its current state, which gets passed to the agent (our decision-making code). The agent selects an action, which drives the motors, which changes the robot's state, generating a new observation. Repeat. During training, the environment also produces a numerical reward at each time step so the agent knows how well it's doing. Everything except the agent (e.g., the chassis, wheels, floor, sensors, and motor control code) is considered the environment.

Designing the Reward Function

Designing the reward function is one of the most consequential decisions in any RL project. It directly shapes what behavior the agent learns, and a poorly designed one will produce an agent that technically maximizes rewards while completely missing the point.

Image of Reinforcement Learning for Robotics Part 2: Train a Balance Bot with PPO

For this balance bot, the reward is computed at every 5ms time step and has four components:

Alive bonus - a small constant reward for staying upright each step
Pitch penalty - subtracts a coefficient times the square of the pitch angle, penalizing the robot for leaning away from vertical
Action penalty - subtracts a coefficient times the sum of squared motor commands, discouraging jerky or aggressive motor inputs
Position and yaw penalties - uses privileged information (the robot's absolute XY position and yaw rate from the simulator) to penalize drifting from the origin and spinning in place

That last point is worth noting. Position and yaw data are available in the simulator but wouldn't be on the real robot. This is known as privileged information, and it's perfectly valid to use during training. We just can't rely on it for deployment.

Building the Gymnasium Environment

Farama Gymnasium is the standard interface for RL environments in Python. Most RL algorithms (including the PPO implementation we're using) know how to interact with a Gymnasium-compatible environment, which makes it a natural wrapper for our MuJoCo simulation.

The full environment code is in workspace/software/ep02/balance_bot_env.py. Here's the structure:

The observation space contains four continuous values:

Pitch (radians, estimated via complementary filter from the IMU accelerometer and gyroscope)
Pitch rate (radians/second)
Left wheel angular velocity
Right wheel angular velocity

The action space contains two continuous values between -1 and 1, representing normalized torque commands for the left and right motors.

The three required methods are reset(), step(), and render(). On reset, the robot returns to the origin upright, but with a small randomized angular velocity (giving it a slight push in a random direction). This is an early form of domain randomization and helps the agent learn to recover from different starting conditions rather than memorizing a single trajectory.

On each step(), the environment applies the action to the motors, advances the simulation by one 5ms timestep, computes the reward, and checks whether the episode has terminated (robot tipped past 30°) or been truncated (hit the maximum step limit).

Before moving on to training, it's worth verifying that the environment works correctly. Open workspace/software/ep02/environment_test.ipynb in JupyterLab and run through the cells. This uses Gymnasium's built-in environment checker and then runs a random policy (choosing actions at random for 200 steps) so you can watch the robot flail around in MuJoCo before any learning has happened. It's a useful sanity check, and it also gives you an intuition for just how hard the balancing task is without a trained policy.

PPO: The Training Algorithm

Proximal Policy Optimization (PPO) is one of the most widely used RL algorithms today. Originally developed by OpenAI researchers and published in 2017, it strikes a good balance between sample efficiency, stability, and ease of implementation.

PPO uses an actor-critic architecture. The actor takes in an observation and outputs a probability distribution over actions, rather than a single deterministic action. This allows for exploration over the possible actions while the agent learns. The critic takes the same observation and estimates the expected total future reward from that state, which is used to help update the actor. In our implementation, both are simple three-layer dense neural networks with 16 nodes per hidden layer (small enough to run inference in under 5 ms on the ESP32).

If you would like to dive more into the math behind PPO, check out my RL math blog post series here.

Our PPO implementation is based on CleanRL, a set of clean reference implementations of RL algorithms created by Costa Huang as part of his research at Drexel University. The version we're using has been modified and documented for this project. We will be using the ppo_continuous_action.py version for our project.

A few things worth understanding about how PPO works:

Rollouts - Before updating the networks, the algorithm collects a batch of experience by running the environments for a fixed number of steps. You'll notice the robot occasionally freezes mid-episode in the viewer; that's the rollout ending and the network update beginning.
Clipping - PPO prevents the actor from being updated too aggressively by clipping the ratio of action probabilities under the new policy versus the old one. This keeps updates stable and prevents the kind of catastrophic forgetting that can derail training.
Annealing - the learning rate gradually decreases over the course of training, helping the agent settle on a stable policy rather than continuing to make large updates late in training.

We run four parallel environments during training (only the first is rendered), which improves sample diversity without overwhelming CPU resources. The timestep in the training loop is set to zero so the environments run faster than real time. The 5ms timestep is still baked into the MJCF physics, so the trained policy behaves correctly when we switch back to real-time evaluation.

Curriculum Learning

One of the trickier aspects of training RL agents is reward hacking, where the agent finds some unintended shortcut that technically maximizes the reward without learning the intended behavior. A classic example is detailed in OpenAI’s 2016 paper talking about how their RL agent learned to maximize rewards in a video game by driving the boat in circles (to collect items) rather than completing the course.

Curriculum learning helps to address this by breaking training into phases of increasing difficulty, shaping the agent's behavior gradually rather than throwing the full task at it from the start.

For this project, we use two phases:

Balance only - The position and yaw penalties are set to zero. The agent only needs to stay upright. It's free to wander and spin as it learns the core skill of balancing first.
Balance and hold position - Using the same agent from Phase 1 as the starting point, we reintroduce the position and yaw penalties. Now the agent needs to stay upright and remain near the origin without spinning. You'll see it tip over more frequently early in this phase as it explores the expanded constraint set, but it eventually converges on the intended behavior.

Running Training

Image of Reinforcement Learning for Robotics Part 2: Train a Balance Bot with PPO

Open workspace/software/ep02/train_with_ppo_curriculum.ipynb in JupyterLab. Run through the cells in order. Once training starts, you can monitor progress in TensorBoard by navigating to localhost:6006 in a browser (in the container or on the host). Here's what to watch:

Episodic length - how long each episode lasts before the robot tips or truncates. Should increase as training progresses.
Episodic return - total reward per episode. The most important metric. Should trend upward.
Entropy - how spread out the actor's action distribution is. Starts high (exploration) and should decrease as the agent converges on a more deterministic policy.
Clip fraction - the proportion of updates that were clipped. A few percent is healthy; consistently high values suggest the learning rate may be too aggressive.
Value loss - how well the critic is predicting returns. Should trend toward zero as the critic improves.

After Phase 1 completes, run the evaluation cell to watch the trained agent in real time. You'll see the robot staying upright but wandering freely as expected. After Phase 2, the robot should balance in place with minimal drift.

Image of Reinforcement Learning for Robotics Part 2: Train a Balance Bot with PPO

Inspecting the Actor Network

Image of Reinforcement Learning for Robotics Part 2: Train a Balance Bot with PPO

Once training is complete, the actor network is exported as an ONNX file. You can inspect it using netron.app. Load the file from workspace/software/ep02/runs//actor.onnx. The architecture is straightforward: a four-element input (pitch, pitch rate, left and right wheel velocities), two hidden layers of 16 nodes each with tanh activations, and a two-element output (left and right motor torques). Simple enough to convert to raw C code for deployment on the ESP32.

What's Next

We have an agent that balances well in simulation. In the next episode, we'll take just the actor network and deploy it to the real robot and quickly discover that the real world doesn't behave quite like the simulator. That gap between simulation and reality (sim-to-real or sim2real), and how to close it, is what the next episode is all about.