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



class PIDOptimizer(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.args = args
        self.kp = nn.Parameter(torch.tensor(1.0, dtype=torch.float32)).cuda()
        self.ki = nn.Parameter(torch.tensor(1.0, dtype=torch.float32)).cuda()
        self.kd = nn.Parameter(torch.tensor(1.0, dtype=torch.float32)).cuda()

    def forward(self):
        return [self.kp, self.ki, self.kd]
