神经网络后[np.arange(0, self.batch_size), action]的目的是啥？

Posted 2023-03-12

【中文标题】神经网络后[np.arange(0, self.batch_size), action]的目的是啥？

【英文标题】：What is the purpose of [np.arange(0, self.batch_size), action] after the neural network?神经网络后[np.arange(0, self.batch_size), action]的目的是什么？

【发布时间】：2022-01-24 05:26:18

【问题描述】：

我按照 PyTorch 教程学习强化学习 (TRAIN A MARIO-PLAYING RL AGENT)，但我对以下代码感到困惑：

current_Q = self.net(state, model="online")[np.arange(0, self.batch_size), action] # Q_online(s,a)

神经网络之后的[np.arange(0, self.batch_size), action]的作用是什么？（我知道TD_estimate接受状态和动作，只是在编程方面对此感到困惑）这是什么用法（在self.net之后放一个列表）？

更多相关代码参考教程：

class MarioNet(nn.Module):

def __init__(self, input_dim, output_dim):
    super().__init__()
    c, h, w = input_dim

    if h != 84:
        raise ValueError(f"Expecting input height: 84, got: h")
    if w != 84:
        raise ValueError(f"Expecting input width: 84, got: w")

    self.online = nn.Sequential(
        nn.Conv2d(in_channels=c, out_channels=32, kernel_size=8, stride=4),
        nn.ReLU(),
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2),
        nn.ReLU(),
        nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1),
        nn.ReLU(),
        nn.Flatten(),
        nn.Linear(3136, 512),
        nn.ReLU(),
        nn.Linear(512, output_dim),
    )

    self.target = copy.deepcopy(self.online)

    # Q_target parameters are frozen.
    for p in self.target.parameters():
        p.requires_grad = False

def forward(self, input, model):
    if model == "online":
        return self.online(input)
    elif model == "target":
        return self.target(input)

self.net：

self.net = MarioNet(self.state_dim, self.action_dim).float()

感谢您的帮助！

【参考方案1】：

本质上，这里发生的是网络的输出被切片以获得 Q 表的所需部分。

[np.arange(0, self.batch_size), action] 的索引（有点混乱）索引每个轴。因此，对于索引为 1 的轴，我们选择由 action 指示的项目。对于索引 0，我们选择 0 到 self.batch_size 之间的所有项目。

如果self.batch_size与这个数组的0维长度相同，那么这个切片可以简化为[:, action]，这可能是大多数用户更熟悉的。