I LoRA fine-tuned then RL'd Gemma 1B to play the game Takeover in a multi-stage training process. Takeover is a competitive card game with a punishing action economy where your goal is to rid your hand of all your cards as quickly as possible.

Takeover is a niche game so I had to implement the game, implement a JAX GPU kernel to model game state transitions and perform legality checks, generate a training corpus by randomly sampling legal moves, train an RNN to produce these legal moves, RL this RNN to produce good moves, fine-tune Gemma 1B on this RNN's outputs, then finally RL Gemma 1B until it achieved a 50% win-rate against the teacher RNN.

This write-up captures the steps involved in the process to produce the strongest Takeover LLM possible.

Overview

Both models—the RNN (henceforth the Teacher) and Gemma 1B (henceforth Gemma)—achieve an 85% win-rate against an opponent who plays a random legal move, a 64% win-rate against an opponent who prioritizes the most aggressive legal move, and a 50% win-rate against one another.

The Takeover interface I implemented looks like so:

The game board visible here represents the game state, which is manipulated by the players as they perform actions in the game. The AI move coach visible here shows sampled moves from the Teacher. The objective of Takeover is to use card special abilities to play "formations" and eventually reduce your hand to zero cards. The number of cards in a player's hand typically decreases turn-over-turn.

The goal of this project is to demonstrate that I can identify a novel problem, develop an environment that the problem can be explored in, develop simple—fast—models to generate comprehensive training data, use this comprehensive training data to teach an LLM to effectively solve the problem, and apply the LLM's solution in a practical way so an end-user can benefit from the LLM's solution to the problem. I achieved this goal.

The greatest difficulties I had in implementing this project were: producing an environment that could rapidly generate thousands of rollouts on my local hardware for training purposes; and identifying an architecture for the "fast" non-LLM model that was performant enough to improve at Takeover during RL.

When implementing this project, I was able to better familiarize myself with RNNs and LSTM networks in particular; though I ended up with a simpler RNN that was fast and good enough for play. I also was able to learn about how to implement a GPU kernel (JAX-compiled) that encoded Takeover's game logic, allowing rollouts to happen on the same device that training was taking place on.

Architecture and Training

The Teacher accepts a 473-wide tensor representing the game state, embeds each input, passes through a few 64-wide feed-forward layers, then a few 64-wide recurrent layers. The output is 14 semantic policy heads relating to various game actions and a value head for RL.

Gemma is fine-tuned/RL'd with a rank-16 LoRA configuration.

The high-level training steps were:

1. Train the small model on randomly sampled valid moves on random games. Continue training until the model has >99.5% success rate for generating a valid move.
2. Self-play RL the small model using PPO, with action validity, cards played, and winning as rewards.
3. Self-play RL the small model using GRPO, with cards played and winning as the reward.
4. SFT Gemma using outputs from the model produced by (3).
5. Self-play RL Gemma using GRPO, with cards played and winning as the reward.

I ran into some issues initially getting my small model to ever output a valid move. I addressed this by allowing a re-sample up to 8 times total. Then I ran into some issues getting my small model to output valid moves on the most complex/most important turn types (Vie and Plan-A). I addressed this by upsampling these data points by 5x during training.

Even with those two changes I still had some valid-move-rate issues so I introduced a PPO step to incentivize simply making a valid move. With PPO, I was able to make a batch that contained the multiple invalid retries, and the correct move and reward them all appropriately.

After the valid move rate was high enough on the small model via PPO, I moved on to GRPO training. I trained using GRPO because I wanted to try it out since it has been popular since the DeepSeek paper. This is where I calculated the 85% win rate against a random-moving opponent for the first time. After performance saturated, I moved on to training Gemma.

I generated some single-turn SFT data that briefly introduces the game, gives the output format, and populates the model response with a move sampled from the small model. I then ran SFT on Gemma with this data for a while. I then ran GRPO on Gemma where an action's reward was +1 for an eventual win (with the Teacher taking over for subsequent moves) and +0 for an eventual loss. I ran this GRPO training until Gemma performance saturated; here I made note that the performance saturated at a 50% win-rate against the small model—it was unable to find an edge.

I ran most training for this project on my maxed 2024 M4 MacBook Pro. However, RL for Gemma took a while and I moved training into Modal for the GRPO step. Here is the timing for a Gemma variant that is 10-10 W/L against the teacher in 20 games, but proved weaker on a 100-game evaluation run.

This Gemma, trained locally, reached a 44% win-rate against the teacher over 200 games.

At this point, I mistakenly judged that I had made some errors in my Gemma hyperparameter selection. I restarted the training process a few times, trying out various ideas. Eventually I realized that I probably just wasn't training for long enough. This is when I moved training to Modal, where I have some credits.

At this point, I instructed ChatGPT/Codex to port the training code to Modal, recreate what I had done thus far but more quickly, then run RL until we reached performance parity.

An Aside on Gemma Prompt Formatting

I tried many different prompt formatting strategies for Gemma, each of which produced different rates of legal moves. The most-legal configuration was the shortest prompt, one that solicited JSON from Gemma (as opposed to a bespoke format or XML), that provided an example of a random legal move from the current position (sampled cheaply from the current game state). General guiding principles for writing an SFT context prompt:
1) Don't let an LLM write the prompt for you (at first).
2) Don't use some bespoke XML format if you can help it—most models are trained to wield JSON well at this point.
3) DO use an LLM to simplify your prompt after you have represented the problem cleanly, it will help align the prompt more with LLMese.

Results

Overall, the Teacher and Gemma dominate the random opponents. Gemma and the Teacher are 50-50 on a 100-game sample. Now I have an opponent to practice against to prepare me for the next time my girlfriend and I play against each other.

Overall, I was pretty frustrated with how much I needed to rely on re-sampling invalid moves from the models—both of them. There are particular complexities of the game state representation that make legal-move-making difficult but I think different architecture selections could improve the Teacher's legality rate.

Future Work

The codebase that produced these results is quite chaotic. I have training and inference done in MLX, JAX, and PyTorch. I have multiple scripts strung together that produce the final models. I have some benchmarking work done to see what the best-performing configurations are but I need to measure this performance both on my Mac as well as on a dedicated GPU. Ideally, I can come to a common representation layer because I know that MLX will be most performant on my Mac and JAX on the remote GPU. Additionally, I need to do a trade study for other libraries/services—especially for LLM training. For example, PrimeRL and Unsloth. TL;DR: Consolidate my training/inference stack.

For simplicity, each of the models that I trained during this project is memoryless. You ought to be able to card count in this game for improved performance, but I figured that the high-order-performance-bit would be simply playing the game at all. I am interested in seeing how a model with the full game state history performs against these models. TL;DR: Train a model that card counts.

Move validity is still pretty low across both the Teacher and Gemma. I have a parameter sweep experiment that reaches 100%@4 for Gemma and 100%@3 for the Teacher. Re-running the full training pipeline with the best variants of each model for each sub-task will show if we can improve total inference time even further by reducing the number of retries necessary. TL;DR: Train a model that relies less on the harness for move validity checking.

Conclusion

I had fun building this out. This is one of the larger-complexity personal projects I've worked on. Certainly one of the furthest out of my previous wheelhouse. I am interested in doing something similar for real-time systems in the future.

I can always be contacted at <contact (at) thornewolf (dot) com>. Talk soon!