import torch
import torch.nn as nn

class PINN_IK(nn.Module):
    def __init__(self, input_dim=3, output_dim=9): # 输出改为 9
        super(PINN_IK, self).__init__()
        # 增加网络深度和宽度，增强拟合能力
        self.net = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.LeakyReLU(0.1), # LeakyReLU 防止梯度消失
            nn.Linear(256, 512),
            nn.LeakyReLU(0.1),
            nn.Linear(512, 512),
            nn.LeakyReLU(0.1),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.1),
            nn.Linear(256, output_dim)
        )

    def forward(self, x):
        return self.net(x)