def __init__()

in agents/model.py [0:0]

• 30 lines of code
• 1 McCabe index (conditional complexity)

    def __init__(self,
                 state_dim,
                 action_dim,
                 device,
                 h_dim=256,
                 t_dim=16):

        super(MLP_Unet, self).__init__()
        self.device = device

        self.action_mlp = nn.Sequential(
            nn.Linear(action_dim, h_dim),
            nn.Mish())
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(t_dim),
            nn.Linear(t_dim, t_dim),
            nn.Mish(),
        )
        self.state_mlp = nn.Sequential(
            nn.Linear(state_dim, state_dim),
            nn.Mish(),
        )

        self.down_1 = EasyBlock(h_dim, t_dim, state_dim)
        self.down_2 = nn.Sequential(nn.Linear(h_dim, h_dim // 2), nn.Mish())
        self.down_3 = EasyBlock(h_dim // 2, t_dim, state_dim)
        self.down_4 = nn.Sequential(nn.Linear(h_dim // 2, h_dim // 4), nn.Mish())

        self.mid_block = EasyBlock(h_dim // 4, t_dim, state_dim)

        self.up_1 = EasyBlock(h_dim // 4, t_dim, state_dim)
        self.up_2 = nn.Sequential(nn.Linear(h_dim // 4, h_dim // 2), nn.Mish())
        self.up_3 = EasyBlock(h_dim // 2, t_dim, state_dim)
        self.up_4 = nn.Sequential(nn.Linear(h_dim // 2, h_dim), nn.Mish())

        self.final_layer = nn.Linear(h_dim, action_dim)