Reinforcement Learning for Robotics Part 6: AI-Powered RC Balance Bot

2026-09-03 | By ShawnHymel

We've made it to the final episode. Our agent can balance, hold position, and follow velocity and yaw commands. Now it's time to give those commands a proper interface. In this episode, we deploy the trained agent to the real robot and build a web-based remote controller: a Wi-Fi access point running on one core of the ESP32, serving a joystick web page that sends commands over a WebSocket connection to the inference loop running on the other core.

By the end, you'll have an AI-powered balance bot you can drive around from your phone!

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

The Architecture

The ESP32 has two cores, which gives us a clean way to separate concerns:

Core 0: runs the Wi-Fi access point and web server, listens for WebSocket connections from the phone, and updates shared command variables when the joystick moves
Core 1: runs the 5ms inference loop: reads sensors, builds the observation vector (including the latest commands), runs actor_forward(), and drives the motors

The two cores share two global variables, cmd_velocity and cmd_yaw, protected by a FreeRTOS mutex to prevent data races (see my Intro to RTOS video series if you need a refresher). Any time Core 0 receives a new joystick command over the WebSocket, it takes the mutex, updates the globals, and releases it. Every 5ms, Core 1 takes the mutex, reads the globals, and uses the values as part of the observation vector. This is a classic producer-consumer pattern, and it's the same kind of concurrency you'd use in any RTOS-based embedded system.

If you want to go deeper on FreeRTOS concepts like mutexes and semaphores on the ESP32, check out my Introduction to RTOS video series.

The Web Page

The controller interface is a single HTML file stored in program memory on the ESP32. It's deliberately minimal: a status indicator (Connected / Not Connected) and a thumbstick: two concentric circles where the inner circle follows your finger within the outer boundary, just like an analog stick on a game controller. Releasing the stick returns it to center and zeroes out the commands.

The JavaScript handles touch and mouse events, maps the stick position to normalized velocity and yaw values between -1 and 1, and sends them over a WebSocket connection to ws://192.168.4.1/ws as a JSON payload:

Copy Code

{ "velocity": 0.72, "yaw": -0.3 }

The ESP32 parses that JSON using the ArduinoJSON library and calls set_commands() to update the shared globals. When the WebSocket disconnects, the commands are automatically zeroed out so the robot stops.

Feel free to customize the web page (add buttons, sliders, or multiple joysticks).

Required Libraries

You'll need three additional Arduino libraries for this sketch. Install them from the Library Manager:

AsyncTCP: asynchronous TCP driver for the ESP32
ESPAsyncWebServer: web server and WebSocket interface built on AsyncTCP
ArduinoJSON: JSON parsing for incoming WebSocket messages

Note that there are older forks of AsyncTCP floating around in the Library Manager. Make sure you install the actively maintained version by ESP32Async.

The Arduino Firmware

Open workspace/software/06-remote-control/balance_bot_rc/ in the Arduino IDE. The actor.h file contains the trained actor network with 48 nodes per hidden layer. Copy this file from your episode 5 training output. The bala.cpp and bala.h files are unchanged from previous episodes.

Key settings at the top of the sketch:

Copy Code

// Wi-Fi access point
const char* SSID = "BalanceBot";
const char* PASSWORD = "balancebot";
IPAddress IP(192, 168, 4, 1);

// Command scaling (must match training values)
const float VELOCITY_FACTOR = 0.5;   // Max m/s
const float YAW_FACTOR = -2.0;  // Max rad/s (negative to correct motor direction)

The velocity and yaw factors must match the maximum values used during training. If the robot turns in the wrong direction when you push the stick, flip the sign on YAW_FACTOR.

Setup initializes the sensors, loads IMU calibration from NVS, zeros the encoders and motors, and then launches the server task on Core 0 using xTaskCreatePinnedToCore(). The main loop() runs on Core 1 as usual.

The server task (Core 0) configures the access point, registers the WebSocket event handler, serves the index page at /, and calls ws.cleanupClients() every second. Once server.begin() is called, the async library handles incoming connections in the background (no polling required).

The inference loop (Core 1) follows the same structure as Episode 5, with one addition: it calls get_commands() each iteration to read the latest velocity and yaw from the shared globals and appends them to the observation vector:

Copy Code

float obs[6] = {
    pitch,
    pitch_rate,
    left_wheel_vel,
    right_wheel_vel,
    cmd_velocity * VELOCITY_FACTOR,
    cmd_yaw * YAW_FACTOR
};

float action[2];
actor_forward(obs, action);

The rest of the code (e.g., clamping, optional post-processing, motor output, and the 5ms busy-wait) is the same as before.

Connecting and Driving

With the firmware uploaded, lay the robot on its front or back face. On your phone:

Connect to the BalanceBot Wi-Fi network (password: balancebot)
Open a browser and navigate to 192.168.4.1
Wait for the status to show Connected — refresh if needed
Stand the robot upright and start driving

The thumbstick controls forward/backward velocity on the vertical axis and yaw rate on the horizontal axis. Releasing the stick returns the commands to zero, and the robot should settle back into a balanced hold.

Image of Reinforcement Learning for Robotics Part 6: AI-Powered RC Balance Bot

A small pitch offset may be needed if the robot consistently leans in one direction. This can be tuned with the PITCH_OFFSET parameter in the firmware. The other post-processing options (motor boost, deadband, low-pass filter) are available but disabled by default; the domain-randomized agent should handle most real-world variation without them.

Where to Go from Here

This is the end of our series! We saw the whole process: a balance bot that learned to balance, hold position, follow commands, and handle real-world conditions entirely through reinforcement learning in simulation. If you want to keep going, here are some directions worth exploring:

Hyperparameter optimization: rather than guessing at learning rate, network size, and phase ordering, tools like Optuna or Ax can automate the search.
Bipedal and quadrupedal locomotion: the same curriculum learning and domain randomization techniques used here scale directly to legged robots. ETH Zurich's robotics group publishes both papers and code that are worth digging into.
Imitation learning: train a policy by having it mimic a recorded human demonstration, rather than learning from a reward function from scratch. Popular for teaching arm manipulators and legged robots.
Inverse reinforcement learning: infer a reward function by observing demonstrations, rather than designing one by hand. Widely used in autonomous driving and manipulation.
Alternative RL algorithms: PPO is a great default, but Soft Actor-Critic (SAC), TD3, decision transformers, and diffusion policy all have strengths in different settings. Worth experimenting with if PPO isn't converging on a new problem.
Real-world agent updates: capturing data from the real robot and using it to fine-tune the policy could help close the sim-to-real gap further, without relying entirely on domain randomization.

If you get your balance bot working or take it further, please share your progress and tag us (@DigiKey) on social media. Good luck, and happy hacking!