dqn_cnn.py

import torch.nn as nn
import torch.nn.functional as F


class ResidualBlock3D(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv3d(channels, channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv3d(channels, channels, kernel_size=3, padding=1),
        )

    def forward(self, x):
        return F.relu(x + self.block(x), inplace=True)


class DQN_CNN(nn.Module):
    def __init__(self, in_channels=16, num_actions=6):
        super().__init__()

        self.stem = nn.Sequential(
            nn.Conv3d(in_channels, 64, kernel_size=2),
            nn.ReLU(),
        )

        self.res_blocks = nn.Sequential(
            ResidualBlock3D(64),
            ResidualBlock3D(64),
            ResidualBlock3D(64),
        )

        self.global_pool = nn.AdaptiveAvgPool3d(1)

        self.norm = nn.LayerNorm(64)

        self.value_head = nn.Sequential(
            nn.Linear(64, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, 1),
        )

        self.adv_head = nn.Sequential(
            nn.Linear(64, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, num_actions),
        )

    def forward(self, x):
        x = self.stem(x)
        x = self.res_blocks(x)

        x = self.global_pool(x)
        x = x.view(x.size(0), -1)

        x = self.norm(x)

        value = self.value_head(x)
        adv = self.adv_head(x)

        q_values = value + adv - adv.mean(dim=1, keepdim=True)
        return q_values