How to build a Neural Network to approximate the Q-function?

Question

I am learning reinforcement learning with Q-learning using online resources, like blog posts, youtube videos, and books. At this point, I have learned the underpinning concepts of reinforcement learning and how to update the q values using a lookup table.

Now, I want to create a neural network to replace the lookup table and approximate the Q-function, but I am not sure how to design the neural network. What would be the architecture for my neural network? What are the inputs and outputs?

Here are the two options I can think of.

1. The input of the neural network is $(s_i, a_i)$ and the output is $Q(s_i,a_i)$

2. The input is $(s_i)$ and the output is a vector $[Q(s_i,a_1), Q(s_i,a_2), \dots, Q(s_i,a_N)]$

Is there any other alternative architecture?

Also, how to reason about which model would be logically better?

Answer 1

I had the same question when I first learned RL. The architectural design may depend on the task you're considering. Since you're moving from a tabular Q-learning to function approximation, I'm suspecting that you are considering a relatively small action space; in this case, you should use option 2 where the input is the state and the number of output nodes matches the number of available actions. One main reason for this choice: when exploiting the learnt function (i.e. policy is $\arg \max_a Q(s,a)$), for option (1) you would have to repeatedly run a forward pass for all available actions, whereas option (2) requires a single forward pass and run arg max either through numpy or the framework you're using.