Homec4science
Diffusion QPFit (master) Browse

QPFit/qpfit.py
3bfb25b5e70amaster

Last Change
Raw File

qpfit.py
View Options

from qpth.qp import QPFunction
import torch
import torch.nn as nn
from torch.autograd import Function, Variable
from torch.nn.parameter import Parameter
import torch.nn.functional as F
from torch.utils.data import Dataset
import scipy.io as sio
# Class to learn a QP from data
#
# in_features - dimension of the parameter (state)
# out_features - dimension of the output (control input)
# nVar - dimension of the QP to learn
#
# Given a parameter x in R^{nBatch x in_features}, the output of the QP is given by
#
# f = linmap_f(x)
# z = min 0.5 z'*(L'*L + eps*I)*z + f'*z s.t. z >= 0
# out = linmap_out(z)
#
# Parameters learned: L, linmap_f, linmap_out
# Note that linmap is affine
class QPNet(nn.Module):
def __init__(self, nVar, in_features, out_features, eps=1e-4, maxIter=5, umax=1):
super().__init__()
self.eps = eps
self.nVar = nVar
self.in_features = in_features
self.out_features = out_features
self.maxIter = maxIter
self.umax = umax
self.L = Parameter(torch.randn(nVar,nVar))
self.f = nn.Linear(in_features=in_features, out_features=nVar)
self.OutMap = nn.Linear(in_features=nVar, out_features=out_features)
# Gz <= h => z >= 0
self.G = -torch.eye(nVar)
self.h = torch.zeros(nVar)
self.I = torch.eye(self.nVar)
self.e = Variable(torch.Tensor())
def to(self, device):
new_self = super(QPNet, self).to(device)
new_self.G = new_self.G.to(device)
new_self.h = new_self.h.to(device)
new_self.I = new_self.I.to(device)
new_self.e = new_self.e.to(device)
return new_self
def forward(self, x):
Q = self.L.t().mm(self.L) + 1e-3 * self.I
x = QPFunction(verbose=-1, eps=self.eps, maxIter=self.maxIter)(Q, self.f(x), self.G, self.h, self.e, self.e)
return torch.clamp(self.OutMap(x), -self.umax, self.umax)
def prepQP(self):
self.Q = self.L.t().mm(self.L) + 1e-3 * torch.eye(self.nVar)
self.e = Variable(torch.Tensor())
def evalQP(self, x):
return QPFunction(verbose=-1, eps=self.eps, maxIter=self.maxIter)(self.Q, self.f(x), self.G, self.h, self.e, self.e)
def save(self, datafile):
def get_numpy(param):
return param.detach().cpu().numpy()
dat = {'F':get_numpy(self.f.weight), 'f':get_numpy(self.f.bias),
'L':get_numpy(self.L),
'O':get_numpy(self.OutMap.weight), 'o':get_numpy(self.OutMap.bias),
'maxIter':self.maxIter,
'type':'QPNet'}
sio.savemat('data/' + datafile + '_fit.mat', dat)
class QPDataset(Dataset):
def __init__(self, mat_file):
"""
Args:
mat_file (string): Name of
"""
mat = sio.loadmat('data/' + mat_file + '.mat')
self.x = torch.tensor(mat['X'].T, requires_grad=False).type(torch.float32)
self.y = torch.tensor(mat['U'].T, requires_grad=False).type(torch.float32)
print("Number of samples")
print(self.x.shape)
print(self.y.shape)
q = mat['U'][0]
cnt = sum(1 for p in q if (p < 1.0 and p > -1.0))
print("Percentage of non-saturated items %.4f%%" % (100.0 * cnt / len(self.y)))
def __len__(self):
return self.x.shape[0]
def __getitem__(self, idx):
return self.x[idx,:], self.y[idx,:]
# Fixed-iterations of the projected gradient method
#
# in_features - dimension of the parameter (state)
# out_features - dimension of the output (control input)
# nVar - dimension of the QP to learn
#
# f = linmap_f(x)
# z = min 0.5 z'*(L'*L + eps*I)*z + f'*z s.t. z >= 0
# out = linmap_out(z)
#
# z+ = relu((L'*L + eps*I)*z + f)
#
# Parameters learned: L, linmap_f, linmap_out
#
class GradNet(nn.Module):
def __init__(self, nVar, in_features, out_features, eps=1e-4, numItr=5, umax=1, device='cpu'):
super().__init__()
self.eps = eps
self.nVar = nVar
self.in_features = in_features
self.out_features = out_features
self.numItr = numItr
self.umax = umax
self.device = device
# self.L = Parameter(torch.randn(nVar,nVar))
self.f = nn.Linear(in_features=in_features, out_features=nVar)
self.OutMap = nn.Linear(in_features=nVar, out_features=out_features)
self.QMap = nn.Linear(in_features=nVar, out_features=nVar, bias=False)
def forward(self, x):
# Q = self.L.t().mm(self.L) + 1e-3 * torch.eye(self.nVar)
f = self.f(x)
z = torch.zeros(x.shape[0], self.nVar).to(self.device)
for i in range(0,self.numItr):
z = self.QMap(z) + f
z = torch.relu(z)
return torch.clamp(self.OutMap(z), -self.umax, self.umax)
# def prepQP(self):
# self.Q = self.L.t().mm(self.L) + 1e-3 * torch.eye(self.nVar)
# self.e = Variable(torch.Tensor())
# def evalQP(self, x):
# return QPFunction(verbose=-1, eps=self.eps, maxIter=self.maxIter)(self.Q, self.f(x), self.G, self.h, self.e, self.e)
def save(self, datafile):
def get_numpy(param):
return param.detach().cpu().numpy()
dat = {'F':get_numpy(self.f.weight), 'f':get_numpy(self.f.bias),
'Q':get_numpy(self.QMap.weight),
'O':get_numpy(self.OutMap.weight), 'o':get_numpy(self.OutMap.bias),
'numItr':self.numItr,
'umax':self.umax,
'type':'GradNet'}
sio.savemat('data/' + datafile + '_fit.mat', dat)
class linNet(nn.Module):
def __init__(self, nVar, in_features, out_features, eps=1e-4, depth=5, umax=1, device='cpu'):
super().__init__()
self.eps = eps
self.nVar = nVar
self.in_features = in_features
self.out_features = out_features
self.depth = depth
self.umax = umax
self.device = device
# self.L = Parameter(torch.randn(nVar,nVar))
self.InMap = nn.Linear(in_features=in_features, out_features=nVar)
self.Map = [nn.Linear(in_features=nVar, out_features=nVar) for i in range(self.depth)]
self.OutMap = nn.Linear(in_features=nVar, out_features=out_features)
def forward(self, x):
z = self.InMap(x)
for i in range(0,self.depth-1):
z = self.Map[i](z)
z = torch.relu(z)
return torch.clamp(self.OutMap(z), -self.umax, self.umax)
# def prepQP(self):
# self.Q = self.L.t().mm(self.L) + 1e-3 * torch.eye(self.nVar)
# self.e = Variable(torch.Tensor())
# def evalQP(self, x):
# return QPFunction(verbose=-1, eps=self.eps, maxIter=self.maxIter)(self.Q, self.f(x), self.G, self.h, self.e, self.e)
def save(self, datafile):
def get_numpy(param):
return param.detach().cpu().numpy()
# dat = {'F':get_numpy(self.f.weight), 'f':get_numpy(self.f.bias),
# 'Q':get_numpy(self.QMap.weight),
# 'O':get_numpy(self.OutMap.weight), 'o':get_numpy(self.OutMap.bias),
# 'numItr':self.numItr,
# 'umax':self.umax,
# 'type':'GradNet'}
# sio.savemat('data/' + datafile + '_fit.mat', dat)
c4science · Help