Introduction #
This is my journey of training a reinforcement learning AI model to race in a video game TrackMania. The goal is to implement a model that learns how to drive and complete the track. Bonus points if it can get some competitive lap times.
This is heavily inspired by Yosh’s YouTube channel where he produces these awesome videos, check it out.
Reinforcement Learning 101 #
The main idea behind reinforcement learning (RL) is that there is no prepared training dataset. The agent needs to learn everything by itself via interacting with an environment.
The environment can be a video game, a simulator or even a real world. It is processed to extract observations - this is what an agent “sees”. It can be as simple as a few numbers or as complex as many camera images and 3D point clouds.
Based on observations, an agent needs to output an action. In our case, actions are steering angle and gas pedal. For a robotic arm it would be motor actuation. Action is chosen to maximize the reward.
After an action has been processed an environment outputs a reward. Good actions should lead to big rewards and bad actions to small or even negative (penalties). For a car, good reward can be proportional to driving in the right direction and negative for hitting the wall or a tree.
Reward design is one of the complex problems that is different for each individual use case. I had a situation when an agent did not receive enough reward driving forward, while crashes were heavily penalized. In the end it just decided it is safer and better to stop. Another classical example is an action game where reward was based on not dying - bots just stopped playing altogether. Balancing “carrot” and “stick” is what took a large chunk of time in this project.
Initially, an agent does not know anything at all and just outputs random actions. Over time, some actions produce more rewards than others. For example, progressing along the track a bit and not slamming into the wall right away. Such “good” behavior is reinforced further and further.
Tech Stack #
The game I decided to work with is TrackMania Nations Forever (2008), it is free. Pretty sure I played it as a kid - funny how things come around!
Most RL projects seem to only run on Windows, but I decided that it can be an interesting challenge to do it on Linux. The game itself is emulated via Proton and runs beautifully at 900 fps. First Linux distro was Kubuntu and later I moved to Arch with niri.
For the RL algorithm I picked Soft Actor-Critic (SAC) which is pretty standard for racing. It really helps that it allows to re-using previous episodes for training, making it very sample efficient. Critical for cases where we are limited by the game and not by the compute. The codebase is Python and neural networks are trained using PyTorch. As you will see later, for fast driving we need to know where the car is in relation to the track. Initially, I tried to read this directly from the game’s memory - it worked, but I needed to find a specific memory address every time the game was relaunched. Attempts to find a stable offset chain did not succeed. Eventually I moved to a simpler solution - TMInterface. It is a mod for TrackMania that allows you to write your own plugins that have access to the game state. There I created a simple script that spawns a TCP server which provides position updates. The training code requests position as it needs it.
Interestingly, unless you train on a captured image directly, there is absolutely no need for a GPU. All the training ran on a Ryzen 7 7800X3D CPU in real time without blocking the game. Moving NNs to the GPU did not bring any speedup since they are so small (around 150k parameters total).
Preparations #
This was my first reinforcement learning project and I wanted to learn how it works. The first few weeks I spent studying and implementing multiple algorithms, each one more complex - DQN, REINFORCE, DDPG, SAC. To make sure they worked I used Gymnasium which has multiple simple games instrumented to be easily used for RL.
Sure, I could have used an existing library or asked an LLM to implement the algorithm, but this would definitely rob me of gaining deep understanding myself. That said, using LLMs for learning and troubleshooting was very helpful.
Once I was done with basic RL, it was time to instrument TrackMania as an RL environment. Generally you need the following:
- Image capture (30 fps is enough)
- Car position and speed from TMInterface
- Game pad emulation to submit actions
Surprisingly, the biggest hurdle initially was screen capture. In modern KDE+Wayland it takes hundreds of lines of obscure code just to grab the window contents. Actually I attempted this project a year ago and failed right here - window capture just did not work. Thankfully, a few months later Opus 4.8 was able to implement a working solution and unblock me. $20 well spent!
For unrelated reasons I moved from KDE to niri in the middle of the project and (of course) the screen capture broke. Again, clankers came to the rescue and generated a few more hundred lines of code that I don’t understand at all. But hey, it works.
First TrackMania training #
My first target was completion of A01 - the game’s first track. Notably, it includes a couple of jumps and elevation changes that will be more important than I thought.
Initially my RL setup was simple: “lidar” and speed as observations. Reward was based on speed.
Lidar rays work the following way - for each frame the system casts 19 rays in front of the car. Long ray - lots of room there. Short ray - near the wall. The agent can (hopefully) learn to avoid driving in the direction of short rays. Why not use the image directly? It is possible, but would require a lot more training, so I decided to start with a simpler option. Very few TrackMania projects use image directly.
Actions are steering (-1 to 1) and gas (0 to 1). Yes, there is no brake - TrackMania is a fast game and many tracks don’t require it.
A speed-based reward is controversial, but it was simple and I thought should work. The track is bounded by walls, so going fast - good, hitting walls and going slow - bad.
After a few days of re-writing my gym code to use TrackMania I put it all together and started the first training run. It did not learn anything at all. It didn’t even clear the first turn.
The next few days were spent troubleshooting. LLM agents and I reviewed the code multiple times and didn’t find anything. Opus said that using speed as a reward could never work and I needed a more complicated approach. Sure, I spent a few more days implementing a more complex reward. Started training - still did not learn anything.
On day five, making another desperate code pass I noticed that I had forgotten to update a previous observation variable. RL training uses pairs of observations for training, so it was learning on garbage data. With this one line fix, the car learned to clear the first turn in a few minutes of training. Phew. It's interesting how despite being very powerful, sometimes LLMs miss simple things and steer you towards unnecessary complexity. Initial pure speed reward did work after all.
Unfortunately, lidar observations did not work on this track due to the track sloping up or down. When the car approaches a top of the climb, color-based lidar rays are telling that there is an obstacle in front (the sky). Let’s postpone this track for now.
Switching tracks #
Yosh (from YouTube) created a simple track to iterate on an algorithm. It is ~30 seconds of continuously snaking road, with no elevation changes or straight lines. This was a great way to keep perfecting my lidar-based model.
With almost no code changes, it learned to complete the track! The first time it happened, the training crashed as it didn’t know what to do with the game’s “Well done, here is your lap time” screen. A good problem to have!
As mentioned previously, using speed as the only reward allows the car to complete the track, but it is not good enough for going fast and chasing lap times. Many similar projects settle on using a pre-recorded trajectory and derive observations from it. To start off, I used the trajectory only to calculate the reward - it is proportional to the number of trajectory points you “collect” per time step. This encourages driving fast and in the right direction (pure speed motivates to drive fast into the wall).
From this point on, I also added additional terms to the reward:
- Time penalty where every second a negative reward is applied. This is meant to incentivize an agent to complete the track faster.
- Wall scrubbing penalty to prevent just scrubbing along the wall constantly
Using a trajectory reward, the best lap time was 37 sec, while pure speed reward never went below 40 sec. Gold medal was still 4 seconds faster. Training took multiple hours and I wanted to see if it can be more efficient. TMInterface has a setting to change gameplay speed and that’s what I wanted to try next. After some fiddling, I came to the conclusion that 2x is the best option. At 4x, training was not quite stable. The hypothesis is that the gamepad latency and NN inference do not scale and controls arrive too late. With this speedup and a few minor changes, it was able to set a new lap record of 33.9 after 10 hours of training. It took just about 10 minutes of training from zero to complete the lap for the first time.
Here I hit the limit on how fast it can go with lidar observations. Next step was using trajectory to infer observations instead of lidar. Based on the car position and the trajectory, the following can be computed:
- Car offset to the nearest trajectory point - i.e. how far off-line the car is
- Car offset to a series of further trajectory points (e.g. 5, 10, 25, 50, 100 meters) - helps to know what kind of turns are coming
So it went from 19 floats of lidar rays to just 6 offset floats. Moreover, image capture was not really needed anymore as it actually drove blind. Note that the trajectory was not a fast one and not following the middle of the track. Just a generic trajectory along the track is enough. As long as there is some reasonable reference point, RL can learn how to drive fast based on it.
With trajectory observations training was getting to good times a lot faster. Final time on this track - 31.1 sec and we have the first AI gold medal!
Back to A01 #
Trajectory-based observations do not have the same flaw as lidar, so I was hoping for a quick success. Naturally, it did not happen and a few adjustments had to be made:
- tracking whether the car fell off the track to terminate an episode
- additional slope observation to show if the track is going up or down
Finally, it was robust enough to learn to complete the lap of A01! After a couple of hours of training it got to 33.9 sec. Gold medal here is 25.8 sec, so still some work to do. This track has 2 big jumps and you can clearly see on the training chart that these are the points where the AI fails most often, even after hundreds of training episodes.
Examples of many, many crashes the AI has to go through to learn:
There is another popular observation that I tried to avoid - track progress. It is highly likely that the model will overfit to a specific track, which did not seem very sportsmanlike at first. However, I was out of levers to pull and decided to give it a go and left it to train overnight. Lo and behold - 28.6 sec lap time, good enough for silver medal in the career mode. Seems that the progress observation allows the model to train much longer and keep improving. Without it, training plateaus at around 1000 episodes.
It was great to see that the car learns how to recover from difficult situations, like from 90-degree slide after a jump:
Initially I thought this was it — gold is not achievable. Then, while working on this write-up, I added one more observation almost as an afterthought: the angle between where the car is pointing and the track direction. This one extra float was worth 3 seconds on A01 — 25.0 sec and a gold medal! Interestingly, on Yosh's track it did not improve the best time at all, just made training converge faster.
One of the fast A01 runs:
Closing Thoughts #
This concludes my experiments for now. I was able to train a reinforcement learning model to drive on two very different tracks, complete them and earn gold medals for both.
Some extra ideas I might try in the future:
- Train for (much) longer, like weeks
- Use pure pixels instead of lidar or trajectory and train on a GPU
- Run multiple game instances at the same time and (hopefully) train faster
- Look into world models like DreamerV3
My internal purist is not very happy that I had to use a known trajectory, but this should be expected I think. People do learn the track to get good lap times. It is unfair to expect an AI to drive fast when every upcoming turn is a complete mystery. Nevertheless, I was happy to see that it is still possible for the AI to learn how to drive and complete the lap using simple image lidar and speed as a reward, just not very fast.
P.S. I did not open source the code yet, but might consider it in the future.