Building Little Monsters to Learn Reinforcement Learning

The other day I stumbled across this BEAUTIFUL demo by Aaron Lemke, where he posted a video of a fuzzy little 5-legged “RL creature” with a voice determined by a five-layer neural network overlaid in the top-left corner of his demo. He also had a dialkit-like right panel with a bunch of sliders modulating the voice’s relationship to the environment.

I’d been pretty averse to actually getting into the nitty-gritty of AI research simply because I was afraid of not being smart enough: the math is hard, I don’t really know if I can derive backpropagation from scratch, jargon is confusing, and plus what’s the use if I can just keep making tool-use agents and pretty UIs with nice microinteractions? Ok, enough with the cope - I was just scared.

That afternoon, I set out on a quest to understand RL well enough to fully create a creature of my own!

What is an RL Mon made of?

I started off by reverse-engineering the demo I had seen by shoving it into ChatGPT. So, in the RL creature example:

Environment: The “world” that the creature exists in is a physics simulator by Google DeepMind called MuJoCo. The body is defined as an MJCF file - it tells legs, joints, mass, etc. I find it very interesting that the body is just a “home” for the policy. I’ve been looking into Iain McGilchrist’s work, where he defines the body as something to be experienced - that our flesh is just a home for our soul and something to be understood analogous to the outside world - that we experience the outside world through our flesh, but we are also experiencing our flesh in the same way! This is similar to when I had the question of “What happens when you take the policy trained for a 2-legged creature and apply it to a 4-legged one?” If you create a generalized policy that has been trained with leggedness as a free degree!

Policy: This is the brain - it’s a neural network that takes the environment as input and the next move as output. (joint positions, velocities, body orientation → torques for each joint).

Learning: The monster uses Proximal Policy Optimization (PPO), a standard RL algorithm in robotics. PPO trains an actor-critic architecture consisting of two neural networks: the actor selects the next action (how to move each joint), and the critic estimates the expected future reward of the current state. The actor learns by trial and error, while the critic learns to estimate whether things are improving or worsening. Together they slowly shift policy weights until the creature consistently achieves the objective.

Hopper

To start off, I was tasked by Chat to begin with Hopper. This guy has his stem fixed vertically and has a joint in his foot - the objective is simply to move forward for as far as possible. I started off using MuJoCo, SB3 Zoo, and my PPO algorithm. I ran it 1M timesteps with a constant lr of 0.0003, and it started bouncing!

I also received some graphs back from TensorBoard - here are some of the more useful definitions:

• Reward - is the creature getting better? Should go up.
• Episode Length - is it surviving longer? Should go up.
• Value Loss - is the critic getting less wrong? Should go down.
• Entropy - is the policy becoming less random? Slowly goes down.
• KL - how much did the policy change this update? Should stay small.
• Explained Variance - how good is the critic? Approaches 1.

Walker

Doubling the number of legs now! Walker has lateral movement locked but can now fall forward and backward. I ran 3M timesteps with constant LR, but the figure kept galloping and struggled to push past 3k timesteps before falling. I was confused wtf, so I started looking into seed randomization, and since I was raw-dogging a start, I ran a few with randomized seeds, but it still sucked

run   ep_rew_mean   ep_len   approx_kl   clip_fraction
s1    2450          825      0.25        0.69
s2    1990          894      0.57        0.74
s3    ~2500         n/a      0.18        0.63

Turns out Walker needs ~50M steps + a stable LR, and my SB3 run only ran 3M steps at a constant 3e-4 (my KL exploding). I decided to switch to BRAX because it runs GPU-parallel, whereas SB3 runs on my Mac’s CPU. BRAX also has better adaptive KL (Kullback-Leibler difference measures how much your policy changed in one update), where the LR is adjusted up/down to hold a KL target, whereas SB3’s target_kl just aborts the update once KL is exceeded.

            OLD          NEW
Framework   PyTorch      JAX
Physics     MuJoCo       MJX
RL          SB3/RL-Zoo   BRAX PPO

After wiring up brax, ran for 56m timesteps. Stopped galloping around 15m timesteps

I also decided it might be fun to make a web viewer so I could have a more intuitive interface, so I had cursor compile MuJoCo to WebAssembly (mujoco-js) for the physics, export each policy to ONNX, and ran it in-browser with onnxruntime-web, and rendered the whole thing with three.js (WebGPU, falling back to WebGL) so everything runs serverless. Then I wrapped it in a quick UI with shadcn + Tailwind, where I could toggle between runs, view the policy visualization and feel smart, compare training curves, and customize the colors in the environment! Had a great time showing people once this was finished :)

Ant

Let’s double the legs! Ant survives (stays upright) from the start, but the actual forward gait doesn’t show up until ~16M steps and only looks good around 33–40M (reward 3400 → 4455). Tops out at ~5900 by 90M. High reward variance throughout.

I then added randomly spawning “food orbs” to make the behavior more entertaining to watch, like in Lemke’s demo!

Not cute enough, what about making food only edible with 2 hands?

Here’s a tree for more context

Aw shit, it keeps flipping over

Ant’s health check is height-only, so it stays “healthy” even when it’s on its back, so it’d just flail. Fixed it so that the torso must now be upright to be considered healthy, and trained specifically for flip-over behavior. Self-righting emerges sharply at ~16M steps and plateaus - not super robust, but it’s kind of like a bottle flip where the ant just kicks its feet and prays - I added a “knock over” button on the web viewer; it’s quite fun to play with.

Future Directions

RL POKEMON: Since the actual policies themselves aren’t super large, I think it would be SO FUN to have dedicated “RL-MON” folders where you can train your own monsters under certain constraints (<100m timesteps, no more than 10 limbs, etc.) and carry them around - either locally or maybe on a personal USB drive (shoutout USB.club). Then you could have designated gym machines where you plug your monsters into, and they have to compete to complete tasks, and the best policy wins - and maybe there are categories for certain types of tasks (speed, strength, problem solving, etc.). I actually think this is probably a genius idea and would be a hit in San Francisco - someone help me make this happen.

On a more serious note tho:

• Generalist policies: how can I efficiently train the ability to flip back over, navigate weird terrain, jump, etc.?
• More Monsters: I want to mess around with bendy limbs, different senses, different environmental properties (e.g. more/less gravity), social dynamics, behavior sharing, etc.
• Expressive features: responsive hair and noises

That sums it up! If you want to tinker: github.com/linder0/rl-mon