Train a Balance Bot With Reinforcement Learning

Reinforcement learning (RL) is a subfield within artificial intelligence (AI) that has made dramatic strides in tackling complex tasks. It involves training an agent to accomplish a task through trial and error. The agent is a computer program that observes its environment and decides what action to take next. The rule it uses to make that decision is called its policy.

The agent can also receive a reward from the environment, such as points in a video game, victory in a game of chess, or making a robot take a step without falling over. This reward is typically a numerical value. The agent uses this reward to improve its policy. Over time and with many interactions with the environment, the agent improves its ability to perform the task given to it by the programmer.

While there will always be a place for classical control schemes, such as a proportional-integral-derivative (PID) loop, a linear-quadratic regulator (LQR), or model predictive control (MPC), RL is gaining a lot of traction in the world of robots. It can be tricky to set up, train an agent, and deploy it to a real robot, but it can produce controllers that are remarkably robust to variations in the environment. Rather than trying to mathematically encode each possible state or scenario, you can train an RL agent to handle a wide variety of scenarios, and then the robot should be able to tackle similar situations.

In this guide, I’ll walk you through a conceptual overview of reinforcement learning and then show you how to train your own balance bot using the popular proximal policy optimization (PPO) algorithm. While balancing a two-wheeled robot has long been solved by classical control methods, this guide acts as an introduction to RL concepts and demonstrates how to train an RL agent, which should take only a few hours. You’ll see how the training and deployment process works, preparing you to tackle more complex robots, like Pollen Robotics/Hugging Face’s new bipedal Microduck.

You can accomplish most of this tutorial using just your computer, as we’ll be simulating most of the training. If you want to deploy to real hardware, you’ll need to purchase an M5Stack Bala2-Fire self-balancing robot kit.

A Brief History of RL

Reinforcement learning dates back to the 1950s with Richard Bellman’s research into dynamic programming, Marvin Minsky’s work in computational trial-and-error learning, and Arthur Samuel’s famous checkers-playing program. Most of this early work in RL was influenced by animal-learning research, notably Thorndike’s law of effect and B.F. Skinner’s operant conditioning.

In the decades that followed, most RL research focused on algorithms that learn from experience rather than from a complete mathematical model of the environment. Bellman’s dynamic programming works well when you know exactly how the world behaves and the number of possible situations is manageable, but that’s rarely true. Even games with well-defined rules can have too many positions to calculate exactly. In 1992, IBM’s Gerald Tesauro combined temporal-difference learning with a neural network to create TD-Gammon, which taught itself backgammon through self-play and reached near-world-champion level. The real world is harder still. Unlike board games, moving a robotic arm through space means dealing with enormous amounts of variability, such as sensor noise, motor backlash, someone bumping the arm, and so on.

The deep learning revolution that started in the early 2010s marked a turning point for the RL field as well, ushering in the era of deep reinforcement learning. In deep RL, neural networks can estimate how much future reward to expect from a given situation, they can represent the policy itself, or they can do both. As a result, RL has seen a whole host of rapid successes. Most literature points to 2013 being the turning point of this deep RL revolution, with DeepMind’s DQN learning to play dozens of Atari games by examining pixel data.

In 2016, DeepMind’s AlphaGo defeated Lee Sedol, one of the world’s top Go players, in a five-game match that shocked the AI community and demonstrated the power of RL. Then in 2019, OpenAI Five, another RL program, defeated the professional eSports team OG in the popular video game Dota 2. Deep RL has since spread well beyond games. It’s used to fine-tune large language models like ChatGPT and Claude, and it’s driving rapid progress in robotics.

We can also see how deep RL is being used to make incredible progress in robotics. In 2024, Boston Dynamics announced that it had integrated RL into Spot’s locomotion controller, working alongside the classical MPC it had used for years. That same year, Disney Research published a paper on its BDX droids, which use RL to perform expressive, animator-designed movements that give them lifelike personalities.

Reinforcement learning is proving to be a breakthrough technology used to advance the field of robotics. Let’s turn our attention to how it works in practice.

The Training Loop

Reinforcement learning consists of a relatively straightforward learning loop. As we noted earlier, the agent is the decision-making part of the code. Other parts of the code read from sensors, look at pixel data, etc. and provide an observation to the agent. The agent looks at the information in that observation and chooses an action as an output.

The output of the agent is used to drive whatever we’re hoping to accomplish. This could be controller inputs for a video game character, a movie recommendation, a decision to buy or sell a stock, or, in our case, the speed and direction of each motor on the balance bot.

Note that the agent refers to a very particular piece of code: the decision-making policy, the value estimator (if one is used), and the learning algorithm that updates them. Everything else, including the supporting code that reads sensors and calculates rewards, the motors, the electronics, the robot chassis, and the rest of the world, is considered the environment.

After an action is taken, the environment reacts in some way. For our robot, the motors spin the wheels, which tips the robot in some direction, possibly driving it forward or backward. The supporting code then reads new data from the sensors, which include an inertial measurement unit (IMU) and two motor encoders for this example. These sensor measurements become the next observation that is fed to the agent.

During the learning process, we also provide a reward describing how the environment reacted to the agent’s previous decision. For our robot, we’ll construct a single numerical reward from a combination of the following:

• Alive: every timestep that the robot has not fallen, give a constant reward
• Pitch: penalize the robot (negative reward) proportional to the number of degrees it has leaned from fully upright
• Action: penalize the robot for making abrupt changes to its motor commands, which encourages it to make small adjustments to stay upright
• Position: penalize the robot for moving off the origin
• Yaw rate: penalize the robot for turning about its vertical (Z) axis, as we want the robot to stay in place without turning.

The learning algorithm (we’ll use PPO in this example) uses rewards to automatically update the internal workings of the agent so that it ideally takes actions in the future that maximize the total rewards it collects over time.

When we’re happy with the performance of the agent, we can deploy it to the robot (or simulation of the robot). We usually deploy only the policy portion of the agent (often called the actor in deep RL). PPO’s value estimator (often called the critic) helps the learning algorithm during training, but the robot doesn’t need it to make decisions afterward. Additionally, our deployment code can remove the reward signal, because now we only need the agent to take in an observation and make a decision.

In our case, we will perform learning in a simulated environment. This allows us to easily reset the robot if it falls, run the physics engine faster than real life, and run multiple, parallel environments to speed up the process. When we deploy the agent, we must take special care to ensure that the observation values match up with those seen in simulation. That means any filtering, normalization, or unit conversions applied in simulation must also be applied on the real robot. For example, if we used m/s² units for acceleration in simulation but our real sensors provide g-force, we need to convert them to m/s² before feeding that data to the agent.

Policies are also usually designed to output normalized actions, which in our case means a motor control value between –1 and 1. In simulation, this value maps to a continuous command ranging from “full power backward” to “full power forward.” On the real robot, we scale that same value to a PWM signal for the motor driver. For example, if the policy outputs –0.73 for the left motor and our maximum PWM value is 255, we set the driver to reverse and send a PWM value of 186. It’s also good practice to clip the policy’s output to the –1 to 1 range before scaling, since the raw values can occasionally fall outside it. Whatever mapping you choose, the goal is consistency between simulation and reality. An output of 1.0 should mean the same thing to the real motors as it did to the simulated ones.

With an understanding of how the overall training loop works, let’s turn our attention to the PPO algorithm we’ll be using.

Proximal Policy Optimization

The world of RL has dozens of useful algorithms that have been developed over the years. Choosing the correct one can be an art in itself. For this demo, we’re going to use PPO, which was described in a 2017 academic paper written by OpenAI researchers. It is still used today as one of the most popular RL algorithms for controlling robots and fine-tuning LLMs. We won’t cover the gritty details of PPO, as there is a lot of underlying math involved, but let’s go over the basics.

While a PhD student at Drexel University, Costa Huang created reference implementations for a number of RL algorithms stored at github.com/vwxyzjn/cleanrl. We’ll be using a modified version of his ppo_continuous_action.py code. Feel free to look through that file as we briefly talk about how PPO works.

PPO is a type of actor-critic method that uses two separate neural networks in the agent.

The critic is a neural network that takes in the observation, which represents the state of the environment. An observation may capture the full state or only part of it. Our robot’s sensors, for example, can’t tell whether a wheel is slipping. Recall that the observation is a collection of values from our sensor readings (or other input from the environment, such as pixel data). The critic’s purpose is to try to answer the question, “How good is this state?”

In complex cases like our robot, we don’t know exactly how the environment will respond to any given action. Even if we understand the physics of a balance bot, sensor noise and motor variations add randomness that’s impossible to model exactly. So rather than trying to model the environment, we have the critic learn something simpler: an estimate of the total reward we can expect from a particular state.

The robot spends some time actually interacting with the environment (or simulation in our case). During this time, we keep a log of which actions were performed, the observations at each timestep, and the rewards received from the environment. This data collection process is known as a rollout. We then use this rollout information to train the critic to predict the future cumulative rewards from a given state.

Cumulative rewards are known as the return. In practice, rewards further in the future are discounted, so they count for a bit less than immediate ones. This encourages the agent to prefer good outcomes sooner rather than later and keeps the return from growing endlessly. In the code, this future factor is controlled by the gamma value. We use the standard neural network update algorithm known as back propagation to update the internal parameters in the critic to better predict future returns.

As the robot moves through the environment collecting rollouts, the critic ideally improves in its prediction capabilities. It might start out randomly guessing what return is to be expected from a given state. Over time, however, it learns to predict the total cumulative rewards from any given observation (state). This information is crucial for training the actor.

The actor is the neural network that acts as the policy for the agent. Just like the critic, it takes in an observation as input at each timestep. Instead of trying to predict the value of a given state, it provides the suggested action as output. In some deep RL algorithms, the output might be a set of discrete values. In this iteration of PPO, we’re allowing for a continuous set of actions for each motor. In other words, we want the actor to tell us how much to drive each motor based on the current observation.

The actor’s output is not a single action. Rather, the actor provides (for each motor) a mean, which is paired with a standard deviation to describe a Gaussian (normal) probability distribution. In some PPO implementations, the actor also provides a standard deviation with the mean. In the CleanRL implementation, the actor network calculates the mean from the observation, while the standard deviation is a separate value (one per motor) that the learning algorithm adjusts over time. During training, the agent samples a random action from this distribution at each timestep.

Providing a probability distribution like this offers an additional benefit beyond just working for continuous actions: it allows us to control the exploration-exploitation tradeoff during training. The actor might guess at a mean and provide a very wide standard deviation at the start (i.e. the actor says “I don’t know what action to choose, just pick something at random”). As the agent tries random actions, the actor ideally learns which actions provide good future returns. As time goes on, the mean moves around and begins to center on the actions that provide the best returns. In addition, the standard deviation shrinks so that the actor goes from guessing a range of actions (exploration) to providing a narrow range of known-good options (exploitation).

Like the critic, the actor is also trained using the backpropagation technique. Rather than trying to map an input of an observation to a future total return, the actor attempts to maximize the returns due to the actions chosen. Because we don’t know exactly how the environment operates (and therefore cannot calculate exact returns in advance), the actor relies on the critic’s predictions to understand how good a given action might be.

After the agent takes an action, we can ask, “Did things go better or worse than the critic expected?” This difference is called the advantage. If an action led to a better outcome than predicted, the actor adjusts to make that action more likely in similar situations. If it went worse, the actor makes it less likely.

As a result, we have the actor and critic working together in the learning process: the critic looks at an observation to approximate how much total return the agent will receive from that state (assuming the agent follows the policy given by the actor), and the actor uses the information from the critic to figure out which actions to choose to maximize that total return.

Actor-critic methods are incredibly powerful, as they can learn policies to accomplish complex goals in a wide variety of environments, even when we can’t fully mathematically model those environments. However, because the actor and critic rely on each other to learn, they can be extremely finicky and sensitive to the hyperparameters. These are the settings we choose by hand, such as the learning rate and gamma, as opposed to the parameters inside the networks that are updated automatically.

While a number of modern deep RL algorithms rely on this actor-critic setup, PPO’s key contribution is a way to keep learning stable. After each rollout, PPO reuses the collected data for several rounds of updates. However, that data was gathered using the old version of the policy, so if the policy changes too much, the data no longer reflects how the agent actually behaves, and, as a result, a single bad update can undo a lot of progress. To prevent this, PPO uses clipping to limit its updates so that the policy can’t change too drastically in one step. Think of it as a speed limit on learning: the agent improves in small, safe increments rather than in big, risky leaps.

With a basic understanding of PPO, let’s dive into actually training our robot!

Setting Up the Environment

We will use Google Colab to train the robot, which means you do not need to install anything on your computer for this part. Open a web browser and navigate to colab.research.google.com/github/shawnhymel/make-reinforcement-learning-for-robotics/blob/main/software/make_rl_train_balance_bot.ipynb. Note that we will use MuJoCo as our physics simulator to train the agent.

If you have not used Colab before, it is a free runtime environment for Python that allows you to run snippets of code one chunk at a time in a format known as a Jupyter Notebook. Each chunk is known as a cell. Click on the first cell of code (it starts with # Install required packages) to highlight it. Press Shift+Enter to run the current cell and move to the next cell. Continue to press Shift+Enter to run the cells in sequence.

The first collection of cells installs and imports the required libraries. Some of the libraries are custom-made and come from this repository: github.com/ShawnHymel/make-reinforcement-learning-for-robotics. Feel free to look through the files in that repository. Most notably, it contains:

• mechanical/bala2-fire/meshes: STL files for the robot model (a chassis and two wheels)
• mechanical/bala2-fire/bala2-fire-simplified.xml: the MJCF file that tells MuJoCo the following:
  • The difference between timesteps (5ms for this simulation)
  • How to load the STL files as meshes
  • Where the meshes should be positioned in the simulator’s world frame
  • How the meshes are connected and which actuators to use
  • The physical properties of the meshes (e.g. mass, moment of inertia)
  • Which sensors to use and where they’re placed.
• software/balance_bot_env.py: a custom Gymnasium environment that loads the MuJoCo simulator, specifies how to obtain an observation from the simulator, and controls the robot based on the actions provided by the agent
• software/rl/ppo_trainer.py: a modified version of the CleanRL ppo_continuous_action.py file. I removed the integrations for Weights & Biases and Hugging Face, and I added comments to describe what’s happening in the code.
• software/train_with_ppo.ipynb: the Jupyter Notebook file that you are currently interacting with.

The next set of cells assigns values to a number of training settings and hyperparameters. Here we’ll create 8 separate, parallel environments to train the robot. You won’t be able to see training happening in real time, as Colab does not have a separate viewer to let you watch MuJoCo. Each environment will take 500,000 total timesteps, where each step includes taking an action and getting an observation and reward in turn. The training algorithm will use that data to update the actor and critic’s internal parameters.

For the PPO hyperparameters, the actor and critic are separate neural networks, each with 2 hidden layers, and each hidden layer has 32 nodes. By setting the timestep value to 0 here, we tell the simulation to take steps as fast as possible (i.e. not to wait the full 5ms), which means the training simulation runs faster than real time.

We then have a number of helper functions that are used to create the environments from our balance_bot_env.py library and MJCF file, generate videos of the robot in action (i.e. see how well the policy is performing), and load an agent from a previous stage (if needed).

After that, we perform two simple tests of the environment. The first test loads the simulated environment and records a video of the robot falling. This shows what happens when the robot does not take any actions (no policy). The second test loads the same environment and records a video of what happens with a random policy. In other words, a random action is chosen for each motor at each timestep. The robot will likely scoot or twist in one direction before immediately falling.

Feel free to play the movies in the cells directly or download them from make-reinforcement-learning-for-robotics/software/output in Colab. If all the cells run successfully, then we are ready to train the agent.

The Reward Function

Continue to press Shift+Enter to run each of the cells in the section titled Phase 1: Balance only. We will train our robot to balance using a series of phases, each more difficult than the last. In this first phase, we set the following reward coefficients (we use the alive_bonus default of 1.0):

envs = make_envs(

NUM_ENVS,

pitch_penalty_coef=0.5,

action_penalty_coef=0.01,

position_penalty_coef=0.001,

speed_penalty_coef=20.0,

yaw_penalty_coef=0.1,

)

If you dive into the balance_bot_env.py file, you can find where we construct the reward function:

# Compute the penalties

pitch_penalty = self.pitch_penalty_coef * pitch**2

action_penalty = self.action_penalty_coef * np.sum(action**2)

position_penalty = self.position_penalty_coef * (x_pos**2 + y_pos**2)

speed_penalty = self.speed_penalty_coef * (x_vel**2 + y_vel**2)

yaw_penalty = self.yaw_penalty_coef * abs(yaw_rate)

# Reward function: subtract penalties from the alive bonus

reward = self.alive_bonus - pitch_penalty - action_penalty - position_penalty - \

speed_penalty - yaw_penalty

At each timestep, we use the information from the observation and other world values to construct this reward value. If it’s a positive value, then it’s a reward. If it’s a negative value, then it’s a penalty. Let’s look at each of these values from the last line:

• alive_bonus: A simple scalar value awarded as long as the robot has not fallen down
• pitch_penalty: We want the robot to try to stay at 0° (upright). We create a penalty proportional to the square of the robot’s pitch (how far off it is from upright).
• action_penalty: Recall that the chosen action is simply two values usually between –1 and 1. We want to encourage the robot to make small movements with the motors. 
• position_penalty: We penalize the robot when it moves off the place it started. We use privileged information here, as we can get the location of the robot from the simulation, but we don’t know that position data in the real world.
• speed_penalty: We use privileged information about the velocity of the robot in the world frame to calculate its speed (absolute value). We penalize the robot for moving in any direction.
• yaw_penalty: Similar to the position penalty, we want to encourage the robot to avoid turning (around the Z axis), so we create another penalty here.

Note that each reward and penalty is multiplied by some coefficient value that we can set in code before training starts. These provide another set of hyperparameter knobs we can tune in order to encourage the robot to perform some actions or discourage others.

Building a reward function is still somewhat of an art in the RL world. You use it to shape the robot’s behavior (without programming it directly), and you have to be careful to avoid instances of reward hacking, where the agent finds a way to maximize the reward without achieving the outcome you intended. For example, if our reward only included the alive bonus and pitch penalty, the agent might learn to roll across the floor in order to stay upright, which would earn it a full reward score without doing what we intended. This is why we added the position and yaw penalties.

Once training is done in phase 1, scroll back up to look at the Tensor Board cell. Tensor Board is a visualization toolkit that will automatically log and graph values. It is extremely popular for watching how machine learning models train over time.

We show a TensorBoard cell at the start of each training phase that populates with data over time. If you do not see such data showing up, press the reload button in the top-right corner of the visualization (it looks like a circular arrow).

PPO is fairly complex, with lots of intermediate steps and values to watch for. However, the most important graph to look at is the one labeled charts/episodic_return. This shows the total return of each episode of simulation (before the robot falls or time runs out). This should approach some maximum value and stay there, indicating that the agent found a way to control the robot in order to maximize total rewards. In our case, 10,000 is the maximum total reward that the agent can obtain per episode. In the image shown, the robot performed very well for phase 1.

If the robot appears to be balancing in the evaluation video, you may proceed to phase 2.

Domain Randomization

By the end of phase 1, we should have a robot that balances well in simulation. However, as you might guess, simulation is nothing like the real world. We made many assumptions about the robot and environment, including the battery voltage, how long it takes to execute a command given by the agent, the shape of the wheels, sensor noise, and so on. This is known as the sim-to-real gap, and it is an ongoing area of research in the robotics community.

While there are myriad ways to tackle the sim-to-real gap, one of the most powerful techniques in RL is domain randomization(DR). Rather than trying to model the real world exactly in simulation, which is often impossible, we can vary certain properties about the world and the robot. The agent simply learns to deal with these variations while continuing to focus on maximizing its rewards. So long as the agent achieves the intended goal and the real robot’s properties lie within the DR ranges, the robot should still balance.

A simple DR example that’s easy to visualize is adding random nudges to the robot during simulation. This simulates both someone gently pushing the robot or the robot bumping into an object. The agent must learn to handle these non-ideal situations, thus creating a more robust agent.

Most reinforcement learning algorithms start with a random policy. In our case, that means the motors will spin in a random direction with random speed, likely causing the robot to quickly fall down. The simulation resets if the robot tips past 30°, which means the robot will reset very often initially. Most RL algorithms look for when those random actions produce a little more total reward than before. The algorithm updates the agent to take those actions more often. This continues over time, and this trial-and-error approach to learning eventually settles on actions that maximize total rewards.

Rather than try to randomize a lot of properties and train the agent in a single phase, we’re going to introduce the concept of curriculum learning. This technique involves training the agent in successively harder phases, where we slowly increase the difficulty of the task or introduce more DR into the mix. Without curriculum learning, you might find that your agent simply cannot learn the most difficult version of your intended task, or it settles for some unintended actions through reward hacking.

Curriculum learning allows us to slowly shape the agent’s behavior over time. We start with an easier task and increase the difficulty. In our case, we first get the robot to balance. We use that same agent as the starting point in phase 2, where we introduce some DR elements. In this case, we vary the controls and electronics within the robot: sensor noise, up to a 1ms delay for actions to take effect, motor noise, and motor torque (to simulate various battery charge levels).

Press Shift+Enter to run each of the cells in the section titled Phase 2: Domain randomization for controls. Training for this phase will take another 20–30 minutes, so be patient. Once that’s done, refresh the TensorBoard display and check on the charts/episodic_return plot. Make sure it’s still close to about 10,000 reward points per episode.

NOTE: At the end of each training phase, we save the best-performing agent throughout the training process. Even if the agent performed poorly at the end of training, we use the agent with the best total rewards per episode going into the next phase.

Watch the evaluation video for phase 2. The robot should still balance in place. Feel free to try different reward coefficients or DR values to see how the agent responds after additional training.

Once you’re happy with the performance of the robot in phase 2, you may execute the cells in Phase 3: Domain randomization for the physical world. Here, we introduce more DR properties, including nudging the robot in a random direction every now and then, varying the mass of the robot’s parts (I measured them on a scale, but that’s not always exact), and varying the friction between the wheels and the ground (to simulate different surfaces).

Once again, check the TensorBoard charts and evaluation video to ensure the robot is behaving as intended in simulation (i.e. balancing upright without moving or turning). Once you’re happy with the performance, run the final cells in the Clean up and save model section. You’ve trained your balance bot agent!

Deploy the Agent

These final cells close out the environments and export the agent for you to use on the real robot. Recall that the agent for PPO has two separate neural networks: the critic for approximating the value of the current state, and the actor for suggesting actions based on that state. When we deploy the agent, we no longer need the critic, as it was just used for training. The actor outputs the mean of the possible actions, as the standard deviation is dropped. We assume that the mean is the suggested action, so we use that to control our motors.

The final two cells export the actor as an ONNX file and as a .h header file for C. On the left side of the Colab page, click on the folder icon, which allows you to browse the files available in this instance. Click through to make-reinforcement-learning-for-robotics/software/output. Hover over the actor.h file to see the three dots pop up. Click on the three dots and select download, which downloads actor.h to your local computer.

You are welcome to repeat the process to download the ONNX file, but we won’t need it for this demo. You can upload it to the netron.app website in order to inspect the neural network’s layers.

On your local computer, navigate to the project repository at github.com/ShawnHymel/make-reinforcement-learning-for-robotics. Click Code→Download ZIP. Unzip the repository somewhere on your computer.

Copy the actor.h file that you downloaded and paste it in make-reinforcement-learning-for-robotics/software/balance_bot_arduino, overwriting the existing actor.h.

Download, install, and open the Arduino IDE. Go to File→Preferences→Additional boards manager URLs, and add the following URL: https://static-cdn.m5stack.com/resource/arduino/package_m5stack_index.json.

Click OK and let the board definitions update. Open the library manager and search for “M5Unified.” Install the M5Unified library.

Open the Arduino sketch found in make-reinforcement-learning-for-robotics/software/balance_bot_arduino. This should open the actor.h file that you copied in a separate tab. Feel free to look through that file, as it contains the neural network code automatically generated by our training program.

NOTE: This is a relatively simple neural network, so we perform inference with it directly on one of the ESP32’s cores using floating point math. The ESP32 easily runs it within our 5ms control loop. For more complex neural networks, you might need to quantize or compress the model or consider running it on some type of accelerator.

Head back to the balance_bot_arduino.ino sketch and look through the code. In the loop() function, we read from the sensors, construct an observation vector, and feed that as input to our actor neural network (using the actor_forward() function). The neural network’s outputs are the two action values, one for each motor.

We constrain the actor’s output to be between –1 and 1, and then we convert the action outputs to PWM signals for the motors. If the robot has tipped beyond a certain point, we simply shut off the motors so you can move it to a new location.

At the end of the loop, we wait however long is necessary to meet our 5ms timestep requirement.

NOTE: The Arduino code contains a number of pre- and post-processing features, like filters and deadbands. These offer another method of closing the sim-to-real gap, but they start to dive into more classical controls. Most of them have been disabled in this sketch so you can see how the actor behaves without them. If your robot does not balance well, try adjusting some of the parameters (e.g. ACTION_DEADBAND, ACTION_ALPHA) to see if they help with balancing. If your bot drifts while balancing, try slightly setting PITCH_OFFSET to something like –0.05 or 0.05, depending on the drift direction.

Connect your Bala2-Fire robot to your computer. Select Tools→Board→M5Stack→M5Fire. Select the serial port for your robot in Tools→Port. Click Sketch→Upload. Your sketch should compile and be flashed to the ESP32 onboard the Bala2-Fire.

Disconnect the USB cable from the robot and carefully place the robot on a flat surface. After a moment, it should begin to balance on its own. You can try pushing it slightly with your finger. If you push too hard, the robot will fall down, as the policy or motors cannot keep up with such large force changes.

Going Further

Your balance bot might not be perfect; we’re asking a lot of PPO in such a short amount of training time. Try adding more phases, or adjusting the reward function to encourage better balancing. For example, you might have a fourth phase that adds randomization for dr.ridge_prob and dr.ridge_torque_max_nm to simulate the tire ridges. The simulated robot has basic cylinders for wheels with smooth edges whereas the real robot has rubber tires with ridges. Adding small, random torques to the axles can simulate these ridges. Alternatively, you might adjust the direction of gravity slightly to have the agent learn to balance on a slight slope.

If you’d like to dive into the specifics of setting up PPO and the environment as well as see how to add a remote control to drive the bot around, check out my video series on DigiKey’s YouTube channel:

Hopefully, this helps you get started in the world of reinforcement learning. A balance bot is just the beginning: all the concepts we covered here can be used as a foundation for training much more complex robots, like quadrupeds and bipedal walkers. Good luck, and happy hacking!

---

This article appeared in Make: Volume 99.

Featured photo by Shawn Hymel