mlx-landing-rl

A little rocket that learned to land itself. Custom PPO written from scratch in Apple MLX — no Stable-Baselines, no demos, just trial, error, and a few thousand crashes on the way to a soft touchdown.

GitHub (training code): https://github.com/RanveerSoni98/mlx-landing-rl

At a glance

Env	LunarLander-v3 (Gymnasium + Box2D)
Algo	PPO clipped (ε=0.2), GAE (γ=0.99, λ=0.95), from scratch in MLX
Net	ActorCritic 8 → 128 → 128 → 4 + value (18,309 params)
Training	500k steps (~244 rollouts × 2048), Adam 3e-4, ~10–15 min on Apple Silicon
Solved?	Yes — mean ~215 over 5 sampled episodes (official bar: avg ≥ 200)

State is 8 numbers: x, y, vx, vy, angle, angular velocity, left-leg, right-leg. Actions are 4 engines: 0=noop, 1=left, 2=main, 3=right.

Files

File	What
lunar_ppo.npz (best)	Best checkpoint (by training reward) — use this one
lunar_ppo_final.npz	Last checkpoint, for comparison
config.json	Arch + hyperparams + env spec
inference.py	Standalone 15-line fly script (no repo clone needed)

How to fly it

from huggingface_hub import hf_hub_download
import mlx.core as mx
import mlx.nn as nn
import numpy as np

class ActorCritic(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(8, 128)
        self.fc2 = nn.Linear(128, 128)
        self.actor = nn.Linear(128, 4)
        self.critic = nn.Linear(128, 1)
    def __call__(self, x):
        h = nn.relu(self.fc1(x))
        h = nn.relu(self.fc2(h))
        return self.actor(h), self.critic(h).squeeze(-1)

path = hf_hub_download("RanveerSoni98/mlx-landing-rl", "lunar_ppo.npz")
model = ActorCritic()
model.load_weights(path)
mx.eval(model.parameters())

# then per step: sample (don't argmax — see below!)
x = mx.array(obs[None, :].astype("float32"))
logits, _ = model(x)
mx.eval(logits)
action = int(mx.random.categorical(logits, axis=-1).item())

Or just run inference.py from this repo — it does all of the above.

Important: sample, don't argmax. This pilot lands far better when you sample from the policy than when you take the greedy action. Greedy solves ~20%; sampling solves ~60–80%. Sampling keeps a little wiggle room to recover from wobbles.

Training

Full code + instructions on GitHub: https://github.com/RanveerSoni98/mlx-landing-rl

git clone https://github.com/RanveerSoni98/mlx-landing-rl.git
cd mlx-landing-rl
uv sync
uv run train-lander --timesteps 500000
uv run eval-lander --episodes 10

Results

Measured with the release checkpoint (lunar_ppo.npz), sampling mode:

Run	Mean	Solved (≥200)
5 episodes	~215	~80%
Typical	~150–220	~60–80% (stochastic — every batch of flights varies)
Greedy (argmax), 10 eps	~30	~20% (sampling is the way to fly this one)

LunarLander counts as "solved" at an average ≥ 200. This pilot gets there — most flights stick the landing, a few still crater. It's a rocket.

Reward (env default)

+100–140 for landing, −100 for crashing, shaping for distance/velocity/angle, −0.3 per main-engine burn. Nothing custom — straight LunarLander-v3.

Limitations

• macOS + Apple Silicon only (MLX doesn't run elsewhere).
• MLX .npz weights — no Safetensors/PyTorch export yet.
• Sampling-only pilot: greedy mode is weak (see above).
• Single seed (0), single env — no domain randomization.

Citation

@software{soni2026mlxlandingrl,
  author = {Soni, Ranveer},
  title = {mlx-landing-rl: LunarLander PPO from scratch in Apple MLX},
  year = {2026},
  url = {https://github.com/RanveerSoni98/mlx-landing-rl}
}

License

MIT — see GitHub repo for the full text.