<reponame>NNHieu/cleanrl
from torch import nn
import torch.nn.functional as F
from .deq.mon import mon, splitting as sp

# ALGO LOGIC: initialize agent here:
class QNetwork(nn.Module):
    def __init__(self, env):
        super(QNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Conv2d(4, 32, 8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(3136, 512),
            nn.ReLU(),
            nn.Linear(512, env.single_action_space.n),
        )

    def forward(self, x):
        return self.network(x / 255.0)

MON_DEFAULTS = {
    'alpha': 1.0,
    'tol': 1e-5,
    'max_iter': 50
}

def expand_args(defaults, kwargs):
    d = defaults.copy()
    for k, v in kwargs.items():
        d[k] = v
    return d

class MonDEQConv2d(nn.Module):

    def __init__(self, splittingMethod: sp.MONPeacemanRachford, in_dim=28, in_channels=1, out_channels=32, m=0.1, **kwargs):
        super().__init__()
        n = in_dim
        shp = (n, n)
        # self.pool = 4
        # self.out_dim = out_channels * (n // self.pool) ** 2
        linear_module = mon.MONSingleConv(in_channels, out_channels, shp, m=m)
        nonlin_module = mon.MONBorderReLU(linear_module.pad[0])
        self.mon = splittingMethod(linear_module, nonlin_module, **expand_args(MON_DEFAULTS, kwargs))
        # self.Wout = nn.Linear(self.out_dim, 10)

    def forward(self, x):
        # x = F.pad(x, (1, 1, 1, 1))
        z = self.mon(x)
        # z = F.avg_pool2d(z[-1], self.pool)
        return z[-1]

# ALGO LOGIC: initialize agent here:
class QMon(nn.Module):
    def __init__(self, env):
        super(QMon, self).__init__()
        self.pre_mon = nn.Sequential(
            nn.Conv2d(4, 32, 8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2),
            nn.ReLU(),
        )
        self.monconv = MonDEQConv2d(sp.MONPeacemanRachford, in_dim=9, in_channels=64, out_channels=64)
        self.post_mon = nn.Sequential(
            nn.Flatten(),
            nn.Linear(3136, 512),
            nn.ReLU(),
            nn.Linear(512, env.single_action_space.n),
        )

    def forward(self, x):
        out = self.pre_mon(x / 255.0)
        # print(out.shape)
        out = self.monconv(out)
        # print(out.shape)
        out = out[:, :, 1:-1,1:-1]
        # print(out.shape)
        out = self.post_mon(out)
        return out
from torch import nn
import torch.nn.functional as F
from .deq.mon import mon, splitting as sp

# ALGO LOGIC: initialize agent here:
class QNetwork(nn.Module):
    def __init__(self, env):
        super(QNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Conv2d(4, 32, 8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(3136, 512),
            nn.ReLU(),
            nn.Linear(512, env.single_action_space.n),
        )

    def forward(self, x):
        return self.network(x / 255.0)

MON_DEFAULTS = {
    'alpha': 1.0,
    'tol': 1e-5,
    'max_iter': 50
}

def expand_args(defaults, kwargs):
    d = defaults.copy()
    for k, v in kwargs.items():
        d[k] = v
    return d

class MonDEQConv2d(nn.Module):

    def __init__(self, splittingMethod: sp.MONPeacemanRachford, in_dim=28, in_channels=1, out_channels=32, m=0.1, **kwargs):
        super().__init__()
        n = in_dim
        shp = (n, n)
        # self.pool = 4
        # self.out_dim = out_channels * (n // self.pool) ** 2
        linear_module = mon.MONSingleConv(in_channels, out_channels, shp, m=m)
        nonlin_module = mon.MONBorderReLU(linear_module.pad[0])
        self.mon = splittingMethod(linear_module, nonlin_module, **expand_args(MON_DEFAULTS, kwargs))
        # self.Wout = nn.Linear(self.out_dim, 10)

    def forward(self, x):
        # x = F.pad(x, (1, 1, 1, 1))
        z = self.mon(x)
        # z = F.avg_pool2d(z[-1], self.pool)
        return z[-1]

# ALGO LOGIC: initialize agent here:
class QMon(nn.Module):
    def __init__(self, env):
        super(QMon, self).__init__()
        self.pre_mon = nn.Sequential(
            nn.Conv2d(4, 32, 8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2),
            nn.ReLU(),
        )
        self.monconv = MonDEQConv2d(sp.MONPeacemanRachford, in_dim=9, in_channels=64, out_channels=64)
        self.post_mon = nn.Sequential(
            nn.Flatten(),
            nn.Linear(3136, 512),
            nn.ReLU(),
            nn.Linear(512, env.single_action_space.n),
        )

    def forward(self, x):
        out = self.pre_mon(x / 255.0)
        # print(out.shape)
        out = self.monconv(out)
        # print(out.shape)
        out = out[:, :, 1:-1,1:-1]
        # print(out.shape)
        out = self.post_mon(out)
        return out
# ALGO LOGIC: initialize agent here: # self.pool = 4 # self.out_dim = out_channels * (n // self.pool) ** 2 # self.Wout = nn.Linear(self.out_dim, 10) # x = F.pad(x, (1, 1, 1, 1)) # z = F.avg_pool2d(z[-1], self.pool) # ALGO LOGIC: initialize agent here: # print(out.shape) # print(out.shape) # print(out.shape)