Teaching a Neural Network to Land on the Moon

I wanted to see how far a small neural network could get at a classic control problem with no rules written by hand. The task: land a spacecraft gently on a landing pad on the Moon. The result: after about two minutes of training on a laptop, the agent lands well in 83 out of 100 attempts.

This post walks through the whole thing, from a pilot that crashes every time to one that touches down between the flags. Everything here comes from a single Jupyter notebook that you can run yourself, and every number and image in this post was produced by that notebook.

Here is where we start:

And here is where we end up:

The tools

We use three things:

• Gymnasium, a library full of small simulated worlds (called environments) made for practicing reinforcement learning. One of them is Lunar Lander.
• PyTorch, to build and train a small neural network.
• Deep Q Learning (DQN), the learning method. It is the same idea DeepMind used in 2015 to teach a computer to play Atari games.

You do not need any previous experience with reinforcement learning to follow along.

Setup

Install the libraries. The box2d extra installs the physics engine that Lunar Lander uses. Python 3.10 or newer is recommended.

pip install torch "gymnasium[box2d]" numpy matplotlib pillow

Then the imports. We also fix a random seed. Training involves a lot of randomness (random starting positions, random exploration, random network weights). Fixing the seed means that if you run the code again on the same machine you get exactly the same results. I ran the notebook several times and every number matched to the decimal.

import os
import random
from collections import deque

import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from IPython.display import Image, display
from PIL import Image as PILImage

SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

# The network is tiny, so the CPU is fast enough and gives the most reproducible results.
device = torch.device("cpu")

ENV_ID = "LunarLander-v3"
MODEL_PATH = os.path.join("models", "lunar_lander_dqn.pth")

Meet the Lunar Lander

A lander starts near the top of the screen with a random push. The landing pad is always in the middle, between two yellow flags. The goal is to touch down gently on the pad, standing upright.

Every reinforcement learning problem has the same loop:

1. The agent looks at the state (where it is, how fast it is moving).
2. It picks an action (fire an engine or do nothing).
3. The environment moves forward one small step in time and returns a reward (a number that says how good that step was) and the new state.
4. Repeat until the lander lands, crashes, or flies away. One full attempt is called an episode.

The agent never sees the rules. It only sees numbers going in and a reward coming out, and it has to figure out the rest.

The state is a list of 8 numbers:

Index	Meaning
0	horizontal position (0 is the center of the pad)
1	vertical position (0 is the ground)
2	horizontal speed
3	vertical speed
4	angle (tilt)
5	angular speed (how fast it is spinning)
6	1 if the left leg touches the ground, else 0
7	1 if the right leg touches the ground, else 0

The actions are 4 choices:

Action	Meaning
0	do nothing
1	fire the left side engine
2	fire the main engine (pushes up)
3	fire the right side engine

The reward is designed by the people who made the environment. In short:

• Moving closer to the pad and slowing down earns points. Moving away loses points.
• Each leg touching the ground earns 10 points.
• Firing the main engine costs a little fuel (0.3 points per frame), the side engines cost less.
• Crashing costs 100 points. Coming to rest safely earns 100 points.

An episode that scores 200 or more is considered a good landing.

A pilot with no brain

To see what we are up against, I first let an agent press random buttons. This helper plays one episode with any policy (a function that takes a state and returns an action), grabs every frame, and saves the result as an animated GIF. We use the same helper later for the trained agent.

def record_episode(policy, filename, seed=None):
    # Plays one episode with the given policy and saves it as an animated GIF.
    # policy is any function that takes a state and returns an action.
    env = gym.make(ENV_ID, render_mode="rgb_array")
    state, _ = env.reset(seed=seed)
    frames = [env.render()]
    score = 0.0
    done = False

    while not done:
        action = policy(state)
        state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        score += reward
        frames.append(env.render())
    env.close()

    # Keep every second frame and shrink the images so the GIF stays small.
    images = [PILImage.fromarray(f).resize((300, 200)) for f in frames[::2]]
    images[0].save(filename, save_all=True, append_images=images[1:], duration=40, loop=0)
    return score


env = gym.make(ENV_ID)
env.action_space.seed(SEED)  # makes the random button presses reproducible
random_policy = lambda state: env.action_space.sample()

os.makedirs("media", exist_ok=True)
score = record_episode(random_policy, "media/random_agent.gif", seed=SEED)

The random pilot scored minus 351. It falls out of the sky:

Now let us build one that learns.

Part 1: Training

The big idea behind Deep Q Learning

Imagine you had a magic table that, for any situation and any action, told you the total reward you will collect from now until the end of the episode if you take that action and then keep playing well. That number is called the Q value (Q stands for quality).

With that table, flying would be easy: in every state, look up the Q value of each of the 4 actions and pick the biggest one.

The problem is that the state is made of 8 continuous numbers, so there are infinitely many situations. No table could hold them all. Instead we train a neural network to estimate the Q values:

8 numbers (state)  ->  neural network  ->  4 numbers (one Q value per action)

How does the network learn the right numbers if nobody tells it the answers? It uses a simple rule called the Bellman equation:

Q(state, action) = reward + gamma * (best Q value in the next state)

In words: the value of what I do now equals the reward I get right away, plus the value of the best thing I can do afterwards. gamma (0.99 here) makes future rewards count slightly less than immediate ones.

Every time the agent takes a step, we get a real reward, and we can compare the network's guess with this target. The difference is the error, and we adjust the network to make it smaller. Over thousands of steps, the guesses become good.

DQN adds two tricks that make this stable:

• Replay buffer. Instead of learning only from the latest step, we store past experiences in a big memory and learn from random samples of it. This stops the network from forgetting old lessons and from overreacting to a run of similar steps.
• Target network. We keep a second, slowly updated copy of the network to compute the target in the Bellman equation. Using the same network for both the guess and the target is like chasing your own tail.

Finally, the agent has to explore. At the start it knows nothing, so it acts randomly. We use a number called epsilon: with probability epsilon the agent picks a random action, otherwise it picks the action with the highest Q value. Epsilon starts at 1 (fully random) and shrinks after every episode, so the agent explores early and uses what it learned later.

Step 1: The Q network

A small network with two hidden layers of 128 neurons each. It takes the 8 state numbers and outputs 4 Q values, one per action. ReLU is the activation function that lets the network learn non linear patterns.

class QNetwork(nn.Module):
    def __init__(self, state_size, action_size, hidden_size=128):
        super().__init__()
        self.fc1 = nn.Linear(state_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, action_size)

    def forward(self, state):
        x = F.relu(self.fc1(state))
        x = F.relu(self.fc2(x))
        return self.fc3(x)

Step 2: The replay buffer

This is the agent's memory. Each experience is a tuple of five things: the state, the action taken, the reward received, the next state, and whether the episode ended.

A deque with a maximum length automatically throws away the oldest memories when it is full. When it is time to learn, we pick a random batch of 64 experiences and turn them into PyTorch tensors.

class ReplayBuffer:
    def __init__(self, capacity, batch_size):
        self.memory = deque(maxlen=capacity)
        self.batch_size = batch_size

    def add(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def sample(self):
        batch = random.sample(self.memory, k=self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)

        states = torch.tensor(np.array(states), dtype=torch.float32, device=device)
        actions = torch.tensor(actions, dtype=torch.int64, device=device).unsqueeze(1)
        rewards = torch.tensor(rewards, dtype=torch.float32, device=device).unsqueeze(1)
        next_states = torch.tensor(np.array(next_states), dtype=torch.float32, device=device)
        dones = torch.tensor(dones, dtype=torch.float32, device=device).unsqueeze(1)
        return states, actions, rewards, next_states, dones

    def __len__(self):
        return len(self.memory)

Step 3: The agent

The agent ties everything together. These are its settings, often called hyperparameters. The values are well known to work for Lunar Lander.

Name	Value	What it does
BUFFER_SIZE	100,000	how many experiences the memory can hold
BATCH_SIZE	64	how many experiences we learn from at once
GAMMA	0.99	how much future rewards count
TAU	0.001	how fast the target network follows the main network
LR	0.0005	learning rate, the size of each correction
UPDATE_EVERY	4	learn once every 4 steps

The agent has three main methods:

• act chooses an action using epsilon (random action or best known action).
• step saves an experience to memory and, every 4 steps, triggers learning.
• learn applies the Bellman equation: it computes the target, compares it with the network's guess, and nudges the network to reduce the error. Then it moves the target network a tiny bit towards the main network.

BUFFER_SIZE = 100_000
BATCH_SIZE = 64
GAMMA = 0.99
TAU = 1e-3
LR = 5e-4
UPDATE_EVERY = 4


class DQNAgent:
    def __init__(self, state_size, action_size):
        self.action_size = action_size

        # The main network (the one that learns) and the target network (a slow copy).
        self.q_network = QNetwork(state_size, action_size).to(device)
        self.target_network = QNetwork(state_size, action_size).to(device)
        self.target_network.load_state_dict(self.q_network.state_dict())

        self.optimizer = torch.optim.Adam(self.q_network.parameters(), lr=LR)
        self.memory = ReplayBuffer(BUFFER_SIZE, BATCH_SIZE)
        self.step_count = 0

    def act(self, state, epsilon=0.0):
        # Explore: with probability epsilon, pick a random action.
        if random.random() < epsilon:
            return random.randrange(self.action_size)

        # Exploit: pick the action with the highest predicted Q value.
        state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)
        with torch.no_grad():
            q_values = self.q_network(state)
        return int(q_values.argmax())

    def step(self, state, action, reward, next_state, done):
        self.memory.add(state, action, reward, next_state, done)

        self.step_count += 1
        if self.step_count % UPDATE_EVERY == 0 and len(self.memory) >= BATCH_SIZE:
            self.learn(self.memory.sample())

    def learn(self, batch):
        states, actions, rewards, next_states, dones = batch

        # Target: reward + gamma * best Q value of the next state (zero if the episode ended).
        with torch.no_grad():
            best_next_q = self.target_network(next_states).max(dim=1, keepdim=True)[0]
            target_q = rewards + GAMMA * best_next_q * (1 - dones)

        # Guess: the Q value the network currently predicts for the action that was taken.
        predicted_q = self.q_network(states).gather(1, actions)

        # Make the guess closer to the target.
        loss = F.mse_loss(predicted_q, target_q)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # Move the target network slightly towards the main network.
        for target_param, param in zip(self.target_network.parameters(), self.q_network.parameters()):
            target_param.data.copy_(TAU * param.data + (1 - TAU) * target_param.data)

    def save(self, path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        torch.save(self.q_network.state_dict(), path)

    def load(self, path):
        weights = torch.load(path, map_location=device)
        self.q_network.load_state_dict(weights)
        self.target_network.load_state_dict(weights)

Step 4: The training loop

This is where the agent actually practices. For each episode:

1. Reset the environment to a new random start.
2. Until the episode ends: choose an action, apply it, get the reward and the new state, and let the agent store and learn from that experience.
3. Shrink epsilon a little so the agent relies less on random moves.

The official bar for "solved" is an average score of 200 over the last 100 episodes. My agent crossed it around episode 600. My first version stopped training right there, and the result was disappointing: that pilot landed well only about half the time. A pilot that just barely passes is still shaky, so the final version keeps practicing for 1500 episodes in total.

def train(n_episodes=1500, max_steps=1000, eps_start=1.0, eps_end=0.01, eps_decay=0.995):
    env = gym.make(ENV_ID)
    env.reset(seed=SEED)
    env.action_space.seed(SEED)

    agent = DQNAgent(state_size=env.observation_space.shape[0], action_size=env.action_space.n)

    scores = []
    last_100 = deque(maxlen=100)
    epsilon = eps_start

    for episode in range(1, n_episodes + 1):
        state, _ = env.reset()
        score = 0.0

        for _ in range(max_steps):
            action = agent.act(state, epsilon)
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            agent.step(state, action, reward, next_state, done)
            state = next_state
            score += reward
            if done:
                break

        scores.append(score)
        last_100.append(score)
        epsilon = max(eps_end, eps_decay * epsilon)

        if episode % 100 == 0:
            print(f"Episode {episode:4d}   average score (last 100): {np.mean(last_100):7.2f}   epsilon: {epsilon:.3f}")

    agent.save(MODEL_PATH)
    print(f"Model saved to {MODEL_PATH}")
    env.close()
    return scores


scores = train()

Step 5: Watching it learn

Training took about two minutes on an Apple M series laptop. Here is the progress log:

Episode  100   average score (last 100): -147.07   epsilon: 0.606
Episode  200   average score (last 100):  -93.93   epsilon: 0.367
Episode  300   average score (last 100):  -54.87   epsilon: 0.222
Episode  400   average score (last 100):   37.58   epsilon: 0.135
Episode  500   average score (last 100):  149.30   epsilon: 0.082
Episode  600   average score (last 100):  206.67   epsilon: 0.049
Episode  700   average score (last 100):  212.91   epsilon: 0.030
Episode  800   average score (last 100):  222.46   epsilon: 0.018
Episode  900   average score (last 100):  240.56   epsilon: 0.011
Episode 1000   average score (last 100):  239.09   epsilon: 0.010
Episode 1100   average score (last 100):  230.48   epsilon: 0.010
Episode 1200   average score (last 100):  229.24   epsilon: 0.010
Episode 1300   average score (last 100):  241.63   epsilon: 0.010
Episode 1400   average score (last 100):  234.09   epsilon: 0.010
Episode 1500   average score (last 100):  234.59   epsilon: 0.010

You can read the story of the agent in those numbers. For the first 300 episodes it crashes constantly. Around episode 400 the average climbs above zero, which means it has learned to hover and not crash. By episode 600 it is landing on the pad most of the time.

The chart makes it clearer. Each light dot is one episode, the solid line is the average of the last 100 episodes, and the dashed line marks 200:

Notice how noisy single episodes stay, even late in training. That is normal: every episode starts in a different random position, and epsilon still adds a little randomness.

All of what the agent learned now lives in models/lunar_lander_dqn.pth, a file of about 75 KB.

Part 2: Inference (Auto Landing)

Inference means using the trained network, not training it anymore. Two things change compared to training:

• No learning. We never call step or learn. The weights stay fixed.
• No exploration. Epsilon is 0, so the agent always picks the action it believes is best.

This part only needs the saved model file, so you never have to train again to watch it fly.

Loading the trained pilot

pilot = DQNAgent(state_size=8, action_size=4)
pilot.load(MODEL_PATH)
pilot.q_network.eval()

trained_policy = lambda state: pilot.act(state, epsilon=0.0)

Watching it land

First, the same starting position that the random agent got at the beginning:

score = record_episode(trained_policy, "media/trained_agent.gif", seed=SEED)

The trained pilot scored 248.7:

Then three more starting positions. Each seed gives a different random push at the start:

for seed in [1, 2, 3]:
    filename = f"media/trained_agent_seed_{seed}.gif"
    score = record_episode(trained_policy, filename, seed=seed)

Seed 1: score 262.6	Seed 2: score 247.9	Seed 3: score 244.4

How reliable is it?

A few nice videos are not proof. So I flew 100 landings without recording video (much faster) and counted how many scored at least 200. I ran the random pilot on the same 100 starting positions for comparison.

def evaluate(policy, n_episodes=100, first_seed=1000):
    env = gym.make(ENV_ID)
    results = []
    for i in range(n_episodes):
        state, _ = env.reset(seed=first_seed + i)
        score = 0.0
        done = False
        while not done:
            state, reward, terminated, truncated, _ = env.step(policy(state))
            done = terminated or truncated
            score += reward
        results.append(score)
    env.close()
    return np.array(results)


trained_scores = evaluate(trained_policy)
random_scores = evaluate(random_policy)

Pilot	Average score	Good landings (200 or more)
Random	minus 203.6	0 of 100
Trained	239.0	83 of 100

The two groups do not even overlap much. The random pilot lives in negative territory. The trained pilot sits mostly between 200 and 300, with a short tail of landings that were safe but a bit sloppy (off center, or burning too much fuel).

Wrapping up

We started with a lander that fell out of the sky and ended with one that lands on the pad, without writing a single rule about how to fly. The agent only had:

• 8 numbers describing its situation,
• 4 buttons it could press,
• a reward signal telling it how well it was doing.

Everything else it figured out by trial and error, guided by the Bellman equation and a small neural network.

The lesson that surprised me most: hitting the official "solved" score is not the same as having a reliable pilot. Stopping at 200 gave me a coin flip. Letting it practice longer turned that into 83 good landings out of 100.

Ideas to explore next

• Change the hyperparameters. Try a smaller learning rate, a faster epsilon decay, or a bigger network, and compare the learning curves.
• Train for less or more. Try train(n_episodes=600) and compare how reliable the pilot is. Then try 3000.
• Make it harder. gym.make(ENV_ID, enable_wind=True) adds wind, and gym.make(ENV_ID, continuous=True) gives the agent a throttle instead of on and off buttons (that one needs a different algorithm).
• Try other methods. Double DQN, Dueling DQN and PPO are natural next steps once this one makes sense.

The full notebook, with every step explained, is on GitHub: drzamoramora/LunarLander. Run lunar_lander.ipynb from top to bottom and you will get the same pilot and the same numbers shown here.