full image - Repost: [P] Building an Reinforcement Learning Agent to play The Legend of Zelda (from Reddit.com, [P] Building an Reinforcement Learning Agent to play The Legend of Zelda)
Mining:
Exchanges:
Donations:
A year go I started trying to use PPO to play the original Legend of Zelda, and I was able to train a model to beat the first boss after a few months of work. I wanted to share the project just for show and tell. I'd love to hear feedback and suggestions as this is just a hobby project. I don't do this for a living. The code for that lives in the original-design branch of my Triforce repo. I'm currently tinkering with new designs so the main branch is much less stable.Here's a video of the agent beating the first dungeon, which was trained with 5,000,000+ steps. At 38 seconds, it learned that it's invulnerable at the screen edge, it exploits that to avoid damage from a projectile. At 53 seconds it steps up to avoid damage, even though it takes a -0.06 penalty for moving the wrong way (taking damage would be a larger penalty.) At 55 seconds it walks towards the rock projectile to block it. And so on, lots of little things the model does is easy to miss if you don't know the game inside and out.As a TLDR, here's an early version of my new (single) model. This doesn't make it quite as far, but if you watch closely it's combat is already far better, and is only trained on 320,000 steps (~6% less than the first model's training steps).This is pretty far along from my very first model.Original DesignI got the original project working using stable-baselines's PPO and default neural network (Shared NatureCNN, I believe). SB was great to get started but ultimately stifling. In the new version of the project I've implemented PPO from scratch with torch with my own simple neural network similar to stable-baseline's default. I'm playing with all kinds of changes and designs now that I have more flexibility and control. Here is my rough original design:Overall StrategyMy first pass through this project was basically "imagine playing Zelda with your older sibling telling you where to go and what to do". I give the model an objective vector which points to where I want it to go on the screen (as a bird flies, the agent still had to learn path finding to avoid damage and navigate around the map). This includes either point at the nearest enemy I want it to kill or a NSEW vector if it's supposed to move to the next room.Due a few limitations with stable-baselines (especially around action masking), I ended up training unique models for traversing the overworld vs the dungeon (since they have entirely different tilesets). I also trained a different model for when we have sword beams vs not. In the video above you can see what model is being used onscreen.In my current project I've removed this objective vector as it felt too much like cheating. Instead I give it a one-hot encoded objective (move north to the next room, pickup items, kill enemies, etc). So far it's working quite well without that crutch. The new project also does a much better job of combat even without multiple models to handle beams vs not.Observation/Action SpaceImage - The standard neural network had a really tough time being fed the entire screen. No amount of training seemed to help. I solved this by creating a viewport around Link that keeps him centered. This REALLY helped the model learn.I also had absolutely zero success with stacking frames to give Link a way to see enemy/projectile movement. The model simply never trained with stable-baselines when I implemented frame stacking and I never figured out why. I just added it to my current neural network and it seems to be working...Though my early experiments show that giving it 3 frames (skipping two in between, so frames curr, curr-3, curr-6) doesn't really give us that much better performance. It might if I took away some of the vectors. We'll see.Vectors - Since the model cannot see beyond its little viewport, I gave the model a vector to the closest item, enemy, and projectile onscreen. This made it so the model can shoot enemies across the room outside of its viewport. My new model gives it multiple enemies/items/projectiles and I plan to try to use an attention mechanism as part of the network to see if I can just feed it all of that data.Information - It also gets a couple of one-off datapoints like whether it currently has sword beams. The new model also gives it a "source" room (to help better understand dungeons where we have to backtrack), and a one-hot encoded objective.Action SpaceMy original project just has a few actions, 4 for moving in the cardinal directions and 4 for attacking in each direction (I also added bombs but never spent any time training it). I had an idea to use masking to help speed up training. I.E. if link bumps into a wall, don't let him move in that direction again until he moves elsewhere, as the model would often spend an entire memory buffer running headlong straight into a wall before an update...better to do it once and get a huge negative penalty which is essentially the same result but faster.Unfortunately SB made it really annoying architecturally to pass that info down to the policy layer. I could have hacked it together, but eventually I just reimplemented PPO and my own neural network so I could properly mask actions in the new version. For example, when we start training a fresh model, it cannot attack when there aren't enemies on screen and I can disallow it from leaving certain areas.The new model actually understands splitting swinging the sword short range vs firing sword beams as two different actions, though I haven't yet had a chance to fully train with the split yet.Frameskip/Cooldowns - In the game I don't use a fixed frame skip for actions. Instead I use the internal ram state of game to know when Link is animation locked or not and only allow the agent to take actions when it's actually possible to give meaningful input to the game. This greatly sped up training. We also force movement to be between tiles on the game map. This means that when the agent decides to move it loses control for longer than a player would...a player can make more split second decisions. This made it easier to implement movement rewards though and might be something to clean up in the future.Other interesting detailsPathfinding - To facilitate rewards, the original version of this project used A* to pathfind from link to what he should be doing. Here's a video of it in action. This information wasn't giving to the model directly but instead the agent would only be given the rewards if it exactly followed that path or the transposed version of it. It would also pathfind around enemies and not walk through them.This was a nightmare though. The corner cases were significant, and pushing Link towards enemies but not into them was really tricky. The new verison just uses a wavefront algorithm. I calculate a wave from the tiles we want to get to outwards, then make sure we are following the gradient. Also calculating the A* around enemies every frame (even with caching) was super slow. Wavefront was faster, especially because I give the new model no special rewards for walking around enemies...faster to compute and it has to learn from taking damage or not.Either way, the both the old and new models successfully learned how to pathfind around danger and obstacles, with or without the cheaty objective vector.Rewards - I programmed very dense rewards in both the old and new model. At basically every step, the model is getting rewarded or punished for something. I actually have some ideas I can't wait to try out to make the rewards more sparse. Or maybe we start with dense rewards for the first training, then fine-tune the model with sparser rewards. We'll see.Predicting the Future - Speaking of rewards. One interesting wrinkle is that the agent can do a lot of things that will eventually deal damage but not on that frame. For example, when Link sets a bomb it takes several seconds before it explodes, killing things. This can be a massive reward or penalty since he spent an extremely valuable resource, but may have done massive damage. PPO and other RL propagates rewards backwards, of course, but that spike in reward could land on a weird frame where we took damage or moved in the wrong direction.I probably could have just not solved that problem and let it shake out over time, but instead I used the fact that we are in an emulator to just see what the outcome of every decision is. When planting a bomb, shooting sword beams, etc, we let the game run forward until impact, then rewind time and reward the agent appropriately, continuing on from when we first paused. This greatly speeds up training, even if it's expensive to do this savestate, play forward, restore state.Neural Networks - When I first started this project (knowing very little about ML and RL), I thought most of my time would be tuning the shape of the neural network that we are using. In reality, the default provided by stable-baselines and my eventual reimplemnentation has been enough to make massive progress. Now that I have a solid codebase though, I really want to revisit this. I'd like to see if trying CoordConvs and similar networks might make the viewport unncessary.Less interesting details/thoughtsHyperparameters - Setting the entropy coefficinet way lower helped a TON in training stable models. My new PPO implementation is way less stable than stable-baselines (ha, imagine that), but still converges most of the time.Infinite Rewards - As with all reinforcement learning, if you give some way for the model to get infinite rewards, it will do just that and nothing else. I spent days, or maybe weeks tweaking reward functions to just get it to train and not find a spot on the wall it could hump for infinite rewards.Debugging Rewards - In fact, building a rewards debugger was the only way I made progress in this project. If you are tackling something this big, do that very early.Stable-Retro is pretty great - Couldn't be happier with the clean design for implementing emulation for AI.Torch is Awesome - My early versions heavily used numpy and relied on stable-baselines, with its multiproc parallelization support. It worked great. Moving the project over to torch was night and day though. It gave me so much more flexibility, instant multithreading for matrix operations. I have a pretty beefy computer and I'm almost at the same steps per second as 20 proc stable-retro/numpy.Future IdeasThis has already gone on too long. I have some ideas for future projects, but maybe I'll just make them another post when I actually do them.Special ThanksA special thanks to Brad Flaugher for help with the early version of this, Fiskbit from the Zelda1 speedrunning community for help pulling apart the raw assembly to build this thing, and MatPoliquin for maintaining Stable-Retro.Happy to answer any questions, really I just love nerding out about this stuff.
Social Media Icons