Training a DQN Agent to Collect Bananas in Unity
A DQN agent learns to collect bananas in Unity. Here are the Q-network, the replay buffer, the target network and how it reached +13 in 481 episodes.
On this page
This is the first reinforcement learning (RL) agent I trained from start to finish. Reinforcement learning means the agent is never told the right answer. It tries things, gets a score, and slowly works out what earns more points.
The game is Unity’s Banana Collector: a square room full of yellow and blue bananas. A yellow banana is worth +1 and a blue one is worth −1. The agent doesn’t see a picture of the room. It only gets a list of numbers, and from those numbers it has to figure out which banana is which.
The code is in the repo.
The environment¶
At each step the agent gets 37 numbers. They describe how fast it is moving and what is in front of it, a bit like a set of laser pointers that report what they hit. Then it picks one of four actions:
- Move forward.
- Move backward.
- Turn left.
- Turn right.
The task counts as solved when the average score stays at +13 or higher over 100 episodes in a row. One lucky episode is not enough.
The model¶
The method is DQN, short for deep Q-network. The idea is simple: a small neural network looks at the 37 numbers and gives a score to each of the four actions. The score is its guess of how many points that action will earn in the long run. The agent picks the action with the highest score.
Because the input is already a short list of numbers and not an image, the network can be tiny. It has two hidden layers of 64 units each:
class QNetwork(nn.Module):
def __init__(self, state_size, action_size, seed, fc1_units=64, fc2_units=64):
super().__init__()
self.seed = torch.manual_seed(seed)
self.fc1 = nn.Linear(state_size, fc1_units)
self.fc2 = nn.Linear(fc1_units, fc2_units)
self.fc3 = nn.Linear(fc2_units, action_size)
def forward(self, state):
x = F.relu(self.fc1(state))
x = F.relu(self.fc2(x))
return self.fc3(x)
The settings are:
- Replay buffer: 1e5
- Minibatch: 64
- Gamma: 0.99
- Tau: 1e-3
- Learning rate: 5e-4
- Network update: every 4 steps
Two parts of this setup do most of the work:
- The replay buffer. This is the agent’s memory. It stores past moves and hands back a random handful to learn from. Without it, the network learns from a long run of nearly identical frames and forgets everything else.
- The target network. This is a slow-moving copy of the main network. The main network learns by chasing a target score, and if that target moves on every update, it never settles. The slow copy (tau = 1e-3) keeps the target steady.
Take away either one and training falls apart.
Training¶
cd scripts
python train_agent.py --unity-app Banana.app --target-score 13
Episode 100 Average Score: 1.065
Episode 200 Average Score: 4.18
Episode 300 Average Score: 7.93
Episode 400 Average Score: 11.20
Episode 481 Average Score: 13.07
Environment solved in 481 episodes! Average Score: 13.07

The agent solved the task in 481 episodes. What I like about this curve is how boring it is. It climbs steadily the whole way, with no long flat stretch and no sudden jump. The agent just gets a little better at judging bananas every episode. The next two posts in this series are not this polite.
To watch a trained agent next to an untrained one, run:
python test_agent.py --unity-app Banana.app --checkpoint-file ../checkpoints/checkpoint_481.pth
Possible improvements¶
I kept this project to plain DQN. Three well-known upgrades would be the next things to try:
- Double DQN. Plain DQN tends to be too optimistic about its scores. This fixes that.
- Dueling DQN. It splits the network in two: one part judges how good the situation is, the other judges how much each action matters.
- Prioritized replay. It replays the surprising moments more often than the routine ones.
The repo has the setup instructions and the Unity downloads for each platform:
The dqn and mlagents scripts come from the Udacity deep-reinforcement-learning repo (MIT), with small changes.