Why do we calculate the mean squared error loss to improve the value approximation in Advantage Actor-Critic Algorithm?

Question

class AtariA2C(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(AtariA2C, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
        )

        conv_output_size = self. _get_conv_out(input_shape)

        self.policy = nn.Sequential(
            nn.Linear(conv_output_size, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions),
        )

        self.value = nn.Sequential(
            nn.Linear(conv_output_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1),
        )

    def _get_conv_out(self, shape):
        o = self.conv(T.zeros(1, *shape))
        return int(np.prod(o.shape))

    def forward(self, x):
        x = x.float() / 256
        conv_out = self.conv(x).view(x.size()[0], -1)
        return self.policy(conv_out), self.value(conv_out)

In Maxim Lapan's book Deep Reinforcement Learning Hands-on, after implementing the above network model, it says

The forward pass through the network returns a tuple of two tensors: policy and value. Now we have a large and important function, which takes the batch of environment transitions and returns three tensors: batch of states, batch of actions taken, and batch of Q-values calculated using the formula $$Q(s,a) = \sum_{i=0}^{N-1} \gamma^i r_i + \gamma^N V(s_N)$$ This Q_value will be used in two places: to calculate mean squared error (MSE) loss to improve the value approximation, in the same way as DQN, and to calculate the advantage of the action.

I am very confused about a single thing. How and why do we calculate the mean squared error loss to improve the value approximation in Advantage Actor-Critic Algorithm?

Answer 1

I believe that the author is referring to how the networks are trained in Deep RL. Consider Deep Q-Learning where the $Q(s,a)$ is approximated using a neural network. Then the loss function used to train the network is $$\mathbb{E}[(r + \gamma \max_{a'} Q(s',a') - Q(s,a))^2]\;.$$ Here, $r + \gamma \max_{a'} Q(s',a')$ is your target, what you want your network to aim towards, and $Q(s,a)$ is what your network predicted. (Note that I have left off some details that can be found in the Nature paper for simplicity).

As for actor-critic methods, most popular actor-critic methods will use the value function to 'replace' the action-value function by using the following relationship: $$\mathbb{E}[r + \gamma v_\pi(s')] = Q_\pi(s,a)\;.$$ This relationship can be proved by looking at exercise 3.13 (or somewhere around there) in the Sutton and Barto textbook. This looks like what the author is doing in the textbook you are reading.

Based on what I said at the start regarding how state-action value functions are trained, it is analogous to train a critic network that approximates the value function in the same way.

e1: spelling.

e2: added link to nature paper.