From Lunar Landings to Bioengineering Breakthroughs: Mastering DQNs for Precision and Innovation

Image generated with DALL·E and post-processed on Photopea

Artificial Intelligence (AI) is becoming a growing part of society and might be the newest technological revolution like the Industrial Revolution in 1760–1840 or the internet in 1983. Deep Reinforcement Learning (RL) is a recent and exciting branch of AI that mimics and aims to supersede the way humans explore and learn in an environment. It is especially relevant in engineering where it is used to solve complex problems like optimising medical diagnosis, treatment optimisation and drug discovery (1). One of the earliest techniques of deep RL is Deep Q-Networks (DQN) which merge Q-Learning with deep neural networks (NN).

This tutorial guides you through DQN in simple terms, guiding you through its fundamentals, implementation, and potential uses in solving bioengineering challenges by attempting to solve OpenAI Gym’s “Lunar Lander v2” *(2)* environment by training a DQN agent. We will observe how an Agent with no knowledge of any principles learns to land a rocket-propelled spacecraft in a designated area while consuming minimal fuel and in harsh conditions including wind.

1. Background information

1.1. Reinforcement learning

Reinforcement Learning (RL) is a framework where an agent learns how to make decisions by conducting trial and error in an environment. The agent observes its current state (S) in the environment, takes an action (A) and receives penalties or rewards (R). The agent refines his decision-making process through a Value function (V), which estimates future rewards, and a policy (π), which guides optimal actions. This allows him to refine his behaviour over time to maximise the cumulative reward and take the best possible course of action.

Reinforcement learning diagram ( image by author)

In Reinforcement learning, the agent learns by developing and enhancing an optimal policy that indicates which action is better suited for specific states. This policy is refined by exploring the environment and calculating an assigned value to each action by computing an estimate of immediate rewards and discounted future rewards. The Bellman equation gives this value:

Annotated Bellman equation ( image by author)

The agent has to choose between 2 strategies: a safer one (exploitation) relying on known high-reward actions or a more adventurous one (exploration) leading to potentially better rewards. The ε-greedy method is commonly used to balance exploration and exploitation by taking random actions with probability ε and exploiting its learned Q-values by taking the action with the highest Q-value with probability 1− ε. *(3)(4)*

1.2. Tabular Q-Learning

Q-learning is a popular RL algorithm used to learn an optimal policy. In this method, the Q-function (Q(s, a)) estimates the cumulative reward for performing action a in state s with the Bellman equation and is stored in a table of state-action values. This Q table acts as the agent’s “conscience” and guides its actions. While this method is very effective in small finite environments, it becomes very limiter to use for larger environments and runs into the following issues:

· Memory inefficiency: As the state-action space grows, storing a large Q-table becomes impractical, especially in complex environments.

· Generalization difficulty: Q-tables struggle to handle high-dimensional or continuous state spaces, such as raw pixel inputs in games. For example, in a simple cart-pole environment using a 64-pixel, grey-scale image as input data there will be 25664x64 states.

To address these limitations, instead of using a Q-table, we use a Q-function, which can be approximated using neural networks and has the same utility of mapping state-action pairs to a Q value. NN are excellent function approximators as they can generalise in a scalable way, using experience. *(5)(6)(7)(8)*

Tabular Q-Learning vs Q-Function learning (image by author)

2. Overview of Deep Q-Networks

Deep Q-Networks address the limitations of tabular Q-learning by using a neural network to approximate the Q-function.

2.2. DQN Architecture

The architecture of DQNs is based on three main components:

· Q Network: a neural network that outputs the Q function for all actions given a state as input.

· Target Network: a secondary network that outputs stable Q-value targets and is updated periodically from the Q network. Prevents unstable training due to bootstrapping by using a separate target network, updated periodically with a copy of the Q network and used to calculate the TD error giving Q time to settle from high fluctuations.

· Experience Replay Buffer: a memory buffer that stores past experiences as tuples (s,a,r,s′) and from which random batches are sampled during training. Prevents overfitting of the latest experienced episodes by storing and replaying past experiences to train the neural network on more than just the latest episode and avoids catastrophic forgetting of associated rewards*(7)(4)(10)(11)*

DQN workflow (image by author) (1)(9)

2.2. DQN Operation in Depth

Annotated DQN Pseudocode (annotation by author) (11)

1. Initialisation

· The Q Network is randomly initialised, some states are chosen, and the Q Network generates random actions to bootstrap the data.

· The Q Network is updated with random weights and copied to the target network.

· The Replay Buffer is populated with a few random interactions.

2. ε-greedy exploration

· The current s state is fed into the Q Network that selects an action using the ε-greedy policy, balancing between exploration and exploitation: a = maxa[Q(s,a)]

3. Gathering training data

· The same current state s is fed into the Target Network that selects an ε-greedy action a based on the Q value.

· The next state s’ resulting from the taken action a and its reward r are recorded.

· The tuple (s,a,s’,r) is recorded in the Experience Buffer

4. Q Network & Target Network predicts Q-value

· The DQN model starts to be trained with a random sample taken from the Experience Replay Buffer and inputted into the Q Network and the Target Network.

· The Q Network outputs the Predicted Q value from the action selected while the Target Network outputs the discounted Q value from the next state plus the reward from the sample batch.

Annotated Target Q Value (image by author)

5. Compute loss and train Q Network

· The mean squared loss is computed using the outputs of the Target Q Value and the Predicted Q Value.

Annotated Loss Function (image by author)

· This loss is back propagated into the Q Network and its parameters are updated using gradient descent.

Annotated Update Rule (image by author)

6. Repeat for next time-step

7. Update the Target Network

After a few batches, the Q Network’s parameters are copied into the Target Network *(7)(5)(6)(11)(10)(9)*

3. Solving Lunar Lander

3.1. The environment

Summary of OpenAI’s Lunar Lander enviromnent (image by author) (2)

To better understand how DQNs operate, we can use the practical example of the Lunar Lander — v2 environment by OpenAI in the Gym toolkit *(2)*. In this environment, the agent aims to successfully land a Lunar Lander on a specific landing site of coordinates (0,0) in an optimised way by firing the left, right or main engine at full throttle or by freefalling. To make the environment more realistic, the agent has to work against wind and turbulences.

This environment provides a state space of 8 values including coordinates, velocities, angle and leg contact information detailed in the above figure. The agent can do either of four actions: do nothing, fire the main engine, and fire the left or right orientation engines.

3.2. The Model

Example Neural Network (image by author)

To tackle this problem, a Deep Q Network is used to approximate the Q function and generate an optimal policy. The Neural Network used is illustrated above. It takes the 8 inputs of the environment’s observation space and outputs the action vector of dimension 4. The example code for this tutorial uses a fully connected network model *(13) but other types of NN can be used, like Convolutional NN [(11)](https://paperswithcode.com/paper/human-level-control-through-deep)*.

3.3. Code overview

Here is an overview of the functions for a simple implementation Of the DQN algorithm for the LunarLander-v2 environment. The full code can be found in the notebooks at the end of the tutorial (code found online and modified).

Defining the Environment and Hyperparameters: call the LunarLander environment

env = gym.make('LunarLander-v2')

Building the Q-Network: Design a neural network, specify neurons and connections.

class QNetwork(nn.Module):
    def __init__(self, state_size, action_size, seed):
        ... 

    def forward(self, x):
        ...

        return x

Building the Replay buffer: holds experience tuples (s,a,r,s’) and done. It returns these values in separate tensors in the sample() function.

class ReplayBuffer:
    def __init__(self, buffer_size, batch_size, seed):
        ...

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

    def sample(self):
        ...
        return (states, actions, rewards, next_states, dones)

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

Creating the Agent: The code was modified to take the hyperparameters as input to be able to do a hyperparameter tuning

class DQNAgent:
    def __init__(self, state_size, action_size, seed,
                 BUFFER_SIZE = int(1e5), BATCH_SIZE = 64 , 
                 GAMMA = 0.99, TAU = 1e-3, LR = 1e-4, 
                 UPDATE_EVERY = 4):
        ...

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

    def learn(self, experiences):
        ...

    def update_fixed_network(self, q_network, fixed_network):
        ...

    def act(self, state, eps=0.0):
        ...

    def checkpoint(self, filename):
        ...

Trainingthe Agent: The code was modified to wrap the main loop in a function

def main_loop(env, dqn_agent, EPS_START, EPS_DECAY, EPS_MIN, MAX_EPISODES, MAX_STEPS, PRINT_EVERY, ENV_SOLVED):

    for episode in range(1, MAX_EPISODES + 1):
        state = env.reset()
        score = 0
        for t in range(MAX_STEPS):
            action = dqn_agent.act(state, eps)
            next_state, reward, done, info = env.step(action)
            dqn_agent.step(state, action, reward, next_state, done)
            state = next_state        
            score += reward        
            if done:
                break

            eps = max(eps * EPS_DECAY, EPS_MIN)
            if episode % PRINT_EVERY == 0:
                mean_score = np.mean(scores_window)
                print('\r Progress {}/{}, average score:{:.2f}'.format(episode, MAX_EPISODES, mean_score), end="")
            if score >= ENV_SOLVED:
                mean_score = np.mean(scores_window)
                print('\rEnvironment solved in {} episodes, average score: {:.2f}'.format(episode, mean_score), end="")
                sys.stdout.flush()
                dqn_agent.checkpoint('solved_200.pth')
                break

        scores_window.append(score)
        scores.append(score)

    return scores

4. Results and Analysis

4.1. Proposed agent results

The trained agent’s performance can be evaluated by plotting a learning curve of cumulative rewards over episodes. Visualisations of the agent’s behaviour before and after training demonstrate the improvement provided by the DQN agent and demonstrate the training of the agent over the trials (episodes). While the rewards are at first negative because the agent is not succeeding, as the agent learns the rewards converge towards a reward of 200 suggesting a success.

Training process for the random and the best agent (code at the end)

The Observation space being 8-dimensional, it is not possible to see the full policy in a plot. However, to illustrate the policy, we can plot it in regard to the x and y coordinates. This is the Agent’s output and illustrates the Agent’s decision-making.

Optimal Policy Visualisation for the best agent (code at the end)

The efficacy of the training can also be tested by outputting a gif of the Lunar Lander in play:

DQN agent landing a spaceship safely (code at the end)

4.2. Hyper-Parameter Tuning

To show how hyperparameters affect the learning process the following learning curves are plotted.

The chosen hyperparameters are τ = 0.001, γ = 0.99, α = 0.0001, target update = 4, ε(start) = 1.0, ε(decay) = 0.999, ε(min) = 0.01.

Learning curves comparing the performance of the agent with different parameter values (code at the end)

The discount factor (γ) determines how much an agent values future rewards compared to immediate rewards. A high value prioritises long-term planning, but if that value becomes too high the immediate reward is discarded making the learning process slower and less stable. This is why γ = 0.99 is the best value.

The soft update parameter τ for updating the fixed q network is also a key player in characterising the training. τ=0.001 is the best value as it provides enough flexibility to explore the environment effectively while exploiting good policies.

Finally, the learning rate plays an important part in the training. A smaller learning rate α=0.0001 is better as it allows the agent to update its policy incrementally, leading to more stable and consistent improvements, and higher values might lead to divergence because of overshooting.

5. Bioengineering Applications

DQL can train robotic surgery arms to navigate to specific anatomical targets with high precision, minimising damage to surrounding tissues. Like the Lunar Lander balancing forces to land smoothly, a robotic arm could learn to balance forces and torques to reach a target with precision, and can perform complex surgical tasks, such as suturing and tissue manipulation. This paper discusses an example where the robot (agent) conducts a task on the human body (environment) by moving the probe (actions) to find a feasible scan plane for the sacrum obtaining information with Ultra Sound imaging (State and Reward Observations) *(12)*.

Illustration of DQL use in surgical robotics (recoloured from (12))

Other implementations include but are not limited to Prosthetic Device Control, Drug Delivery Systems (The capsule must reach a target organ (landing zone) while avoiding “turbulence” (blood flow) or “obstacles” (immune system responses)), Medical Imaging Diagnostics (processing Biomedical imaging files).

Conclusion

In conclusion, while DQN has limitations such as high computational demands, sensitivity to hyperparameter tuning, and potential instability during training, it laid the foundation for significant advancements in reinforcement learning. Some variants of DQN such as Double DQN, Dueling DQN, and Prioritized Experience Replay address these shortcomings resulting in a faster convergence time and more stable learning. Despite being considered outdated, DQN models illustrate how Reinforcement Learning can be improved with Neural Networks.

References

1. DQN RL
2. Lunar Lander
3. Sutton, R. S. & Barto, A. G., 1998. “Dynamic Programming,” in Reinforcement Learning: An Introduction. Cambridge(Massachusetts): MIT Press.
4. Solving Gymnasium’s Lunar Lander with Deep Q Learning (DQN)
5. Youtube: DQN Explained
6. Slides: From Tabular Q-Learning to DQN
7. Reinforcement Learning Explained Visually
8. Youtube: DQN theory and Implementation
9. DQN in PyTorch
10. Landing a Space Craft on the Moon Using Deep Reinforcement Learning
11. Mnih, Volodymyr et al. “Human-level control through deep reinforcement learning.” Nature 518 (2015): 529–533.
12. Surgical robotics paper

Code

13. Source

14. Modified version for visuals ( modifications by author)

15. Chat GPT conversation used for visuals