dro 0.0.1__tar.gz

dro-0.0.1/MANIFEST.in

@@ -0,0 +1 @@
+include pydro/README.md

dro-0.0.1/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.1
+Name: dro
+Version: 0.0.1
+Summary: A package of distributionally robust optimization (DRO) methods. Implemented via cvxpy and PyTorch
+Home-page: https://github.com/namkoong-lab/dro
+Author: Jiashuo Liu, Tianyu Wang, Peng Cui, Hongseok Namkoong
+Author-email: liujiashuo77@gmail.com, tw2837@columbia.edu, cuip@tsinghua.edu.cn, namkoong@gsb.columbia.edu
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: POSIX :: Linux
+Requires-Python: >=3
+Description-Content-Type: text/markdown
+License-File: LICENSE
+
+### DRO Package
+
+> <a href="https://ljsthu.github.io">Jiashuo Liu*</a>, <a href="https://wangtianyu61.github.io">Tianyu Wang*</a>, <a href="https://pengcui.thumedialab.com">Peng Cui</a>, <a href="https://hsnamkoong.github.io">Hongseok Namkoong</a>
+
+> Tsinghua University, Columbia University
+
+
+`DRO` is a python package that implements 12 typical DRO methods on linear models (SVM, logistic regression, and linear regression). It is built based on `cvxpy`. Implemented DRO methods include:
+* $f$-DRO
+    * CVaR-DRO
+    * KL-DRO
+    * TV-DRO
+    * Marginal DRO (CVaR)
+* Wasserstein DRO
+    * Wasserstein DRO
+    * Augmented Wasserstein DRO
+    * Regularized Wasserstein DRO
+* MMD-DRO
+* Sinkhorn-DRO
+* Holistic DRO
+* Unified-DRO
+    * $L_2$ cost
+    * $L_{inf}$ cost
+
+Current version only contains linear models. And further version will incorporate neural network implementations via some approximations.

dro-0.0.1/dro/__init__.py

@@ -0,0 +1,7 @@
+# Built on Folktables and the codebase https://github.com/jpgard/subgroup-robustness-grows-on-trees
+__version__ = "0.0.1"
+__author__ = 'Jiashuo Liu, Tianyu Wang, Peng Cui, Hongseok Namkoong'
+__credits__ = 'Peng Cui\'s Group and Hongseok Namkoong\' Group'
+
+
+from .fetch_method import fetch_method

dro-0.0.1/dro/fetch_method.py

@@ -0,0 +1,38 @@
+from .src import * 
+
+
+def fetch_method(method, is_regression, input_dim=77):
+    if method == 'unified_dro_l2':
+        if is_regression == 1 or is_regression == 2:
+            raise NotImplementedError("Unified DRO does not support regression!")
+        else:
+            return MOT_Robust_CLF_L2()
+    if method == 'unified_dro_linf':
+        if is_regression == 1 or is_regression == 2:
+            raise NotImplementedError("Unified DRO does not support regression!")
+        else:
+            return MOT_Robust_CLF_Linf()
+    elif method == 'hr_dro_lr':
+        return HR_DRO_LR(is_regression=is_regression)
+    elif method == 'chi2_dro':
+        return chi2_DRO(input_dim=input_dim, is_regression=is_regression)
+    elif method == 'kl_dro':
+        return KL_DRO(input_dim=input_dim, is_regression=is_regression)
+    elif method == 'tv_dro':
+        return TV_DRO(input_dim=input_dim, is_regression=is_regression)
+    elif method == 'marginal_cvar_dro':
+        return Marginal_CVaR_DRO(input_dim=input_dim, is_regression=is_regression)
+    elif method == 'cvar_dro':
+        return CVaR_DRO(input_dim=input_dim, is_regression=is_regression)
+    elif method == 'wasserstein_dro':
+        return Wasserstein_DRO(input_dim=input_dim, is_regression=is_regression)
+    elif method == 'wasserstein_dro_satisficing':
+        return Wasserstein_DRO_satisficing(input_dim=input_dim, is_regression=is_regression)
+    elif method == 'wasserstein_dro_aug':
+        return Wasserstein_DRO_aug(input_dim=input_dim, is_regression=is_regression)
+    elif method == 'sinkhorn_dro':
+        return Sinkhorn_DRO_Linear(input_dim=input_dim, is_regression=is_regression)
+    elif method == 'mmd_dro':
+        return MMD_DRO(input_dim=input_dim, is_regression=is_regression)
+    else:
+        raise NotImplementedError(f"Algorithm {method} not implemented yet!")

dro-0.0.1/dro/src/HR_DRO_LR.py

@@ -0,0 +1,111 @@
+import cvxpy as cp
+import numpy as np
+from sklearn.metrics import f1_score
+
+class HR_DRO_LR:
+    def __init__(self, r=1.0, alpha=1.0, epsilon=0.5, epsilon_prime=1.0, is_regression=0):
+        self.r = r 
+        self.alpha = alpha 
+        self.epsilon = epsilon 
+        self.epsilon_prime = epsilon_prime
+        self.is_regression = is_regression
+    
+    def update(self, config={}):
+        if 'r' in config.keys():
+            self.r = config["r"]
+        if 'alpha' in config.keys():
+            self.alpha = config["alpha"]
+        if 'epsilon' in config.keys():
+            self.epsilon = config["epsilon"]
+        if 'epsilon_prime' in config.keys():
+            self.epsilon_prime = config["epsilon_prime"]
+        
+    def fit(self, X, Y):
+        T = X.shape[0]
+        theta = cp.Variable(X.shape[1])  # Define the size based on the problem
+        w = cp.Variable(T)
+        lambda_ = cp.Variable(nonneg=True)
+        beta = cp.Variable(nonneg=True)
+        eta = cp.Variable()
+        temp = cp.Variable()
+
+        # Objective
+        objective = cp.Minimize(1/T * cp.sum(w) + lambda_ * (self.r - 1) + beta * self.alpha + eta)
+        
+        if self.is_regression == 1 or self.is_regression == 2:
+            # Constraints
+            constraints = []
+            # Add constraints based on the problem
+            for t in range(T):
+                constraints.append(temp >= cp.abs(theta.T @ X[t] - Y[t]))
+                constraints.append(w[t]>= cp.rel_entr(lambda_, (eta - cp.abs(theta.T @ X[t] - Y[t]) - self.epsilon * cp.norm(theta, 2)) ))
+                constraints.append(w[t]>= cp.rel_entr(lambda_, (eta - temp - self.epsilon_prime * cp.norm(theta, 2)))-beta)
+                constraints.append(eta >= cp.abs(theta.T @ X[t] - Y[t]) + self.epsilon_prime * cp.norm(theta, 2))
+        elif self.is_regression == 0:
+            Y = 2*Y - 1.0
+            # Constraints
+            constraints = []
+            # Add constraints based on the problem
+            constraints.append(eta>=1e-6)
+            for t in range(T):
+                constraints.append(temp <= Y[t]*(theta.T@X[t]))
+                constraints.append(w[t] >= cp.rel_entr(lambda_, eta))
+                constraints.append(w[t]>= cp.rel_entr(lambda_, (eta - 1 + Y[t]*(theta.T @ X[t])-self.epsilon*cp.norm(theta,2)) ))
+                constraints.append(w[t]>= cp.rel_entr(lambda_, (eta - 1 + temp - self.epsilon_prime * cp.norm(theta, 2)))-beta)
+                constraints.append(eta >= 1-Y[t]*(theta.T @ X[t])+self.epsilon*cp.norm(theta,2))
+        else:
+            raise NotImplementedError
+
+        # Problem
+        prob = cp.Problem(objective, constraints)
+
+        # Solve
+        prob.solve(solver=cp.MOSEK,
+            mosek_params={"MSK_DPAR_INTPNT_CO_TOL_REL_GAP": 1e-8},
+            verbose=True)
+
+
+        self.theta = theta.value
+        self.w = w.value 
+        self.lambda_ = lambda_.value 
+        self.beta = beta.value
+        self.eta = eta.value 
+
+        model_params = {}
+        model_params["theta"] = self.theta.reshape(-1).tolist()
+        return model_params
+    
+    def predict(self, X):
+        scores = self.theta.T @ X.T
+        preds = scores.copy()
+        preds[scores >= 0] = 1
+        preds[scores < 0] = 0
+        return preds
+
+    def score(self, X, y):
+        # calculate accuracy of the given test data set
+        predictions = self.predict(X)
+        acc = np.mean([predictions.flatten() == y.flatten()])
+        f1 = f1_score(y, predictions, average='macro')
+        return acc, f1 
+
+
+        
+        
+
+if __name__=="__main__":
+    from sklearn.datasets import make_regression
+    from  sklearn.linear_model import LinearRegression
+    import numpy as np
+
+    sample_size = 1000
+    feature_size = 10
+    X, y = make_regression(n_samples = sample_size, n_features = feature_size, noise = 1, random_state = 42)
+
+    method = HR_DRO_LR()
+    method.fit(X,y)
+    print(method.theta)
+
+    model = LinearRegression(fit_intercept=True)
+    model.fit(X, y)
+    print(model.coef_)

dro-0.0.1/dro/src/MMD_DRO.py

@@ -0,0 +1,116 @@
+from .base import *
+import numpy as np
+import math
+import cvxpy as cp
+from sklearn.metrics.pairwise import rbf_kernel
+from sklearn.metrics import euclidean_distances
+
+"""
+we set (X, Y) as the same scale but it may not be in practice.
+"""
+
+
+class MMD_DRO(base_DRO):
+    def __init__(self, input_dim, is_regression=2):
+        base_DRO.__init__(self, input_dim, is_regression)
+        self.eta = 0.1
+        self.sampling_method = 'bound'
+        self.n_certify_ratio = 1
+    def update(self, config = {}):
+        if 'eta' in config.keys():
+            self.eta = config['eta']
+        if 'sampling_method' in config.keys():
+            assert (config['sampling_method'] in ['bound', 'hull'])
+            self.sampling_method = config['sampling_method']
+        if 'n_certify_ratio' in config.keys():
+            self.n_certify_ratio = config['n_certify_ratio']
+
+    def matrix_decomp(self, K):
+        try:
+            L = np.linalg.cholesky(K)
+        except:
+            # print('warning, K is singular')
+            d, v = np.linalg.eigh(K)     #L == U*diag(d)*U'. the scipy function forces real eigs
+            d[np.where(d < 0)] = 0       # get rid of small eigs
+            L = v @ np.diag(np.sqrt(d))
+        return L
+
+    def medium_heuristic(self, X, Y):
+        if self.is_regression == 1 or self.is_regression == 2:
+            distsqr = euclidean_distances(X, Y, squared = True)
+        else:
+            distsqr = euclidean_distances(X[:,0:-1], Y[:,0:-1], squared = True)
+
+        kernel_width = np.sqrt(0.5 * np.median(distsqr))
+
+        # in sklearn,
+        # kernel is done by K(x, y) = exp(-gamma ||x-y||^2)
+        kernel_gamma = 1.0 / (2 * kernel_width ** 2)
+
+        return kernel_width, kernel_gamma
+
+    def cvx_loss(self, theta, zeta):
+        if self.is_regression == 1 or self.is_regression == 2:
+            loss = (zeta[-1] - theta @ zeta[:-1]) ** 2
+        else:
+            loss = cp.pos(1 - cp.multiply(zeta[-1], theta @ zeta[:-1]))
+        return loss
+
+    def fit(self, X, y):
+        if self.is_regression == 0:
+            y = 2*y-1
+        sample_size, __ = X.shape
+        n_certify = int(self.n_certify_ratio * sample_size)
+
+        theta = cp.Variable(self.input_dim)                  # variable \theta
+
+        # constraint on the decision variable
+
+        # KDRO part
+        a = cp.Variable(sample_size + n_certify)    # variable \alpha
+        f0 = cp.Variable()                       # variable f0
+        
+        # --------------------------------------------------------------------------------
+        # Step 1: generate the sampled support
+        # --------------------------------------------------------------------------------
+        if self.sampling_method == 'bound':         # sample within certain bound
+            # let the samples also live in the intervel I sampled uncertainty w
+            zeta = np.random.uniform(-1, 1, size=[n_certify, self.input_dim + 1])
+        elif self.sampling_method == 'hull':        # sample using convex hull of empirical data
+            # let the samples also live in the intervel I sampled uncertainty w
+            # this is equiv. to really do it in multi dimensions, need to sample coeff. from a simplex
+            if self.is_regression == 1 or self.is_regression == 2:
+                zeta1 = np.random.uniform(np.min(X), np.max(X), size = [n_certify, self.input_dim])
+                zeta2 = np.random.uniform(np.min(y), np.max(y), size = [n_certify, 1])
+            else:
+                zeta1 = np.random.uniform(-1, 1, size = [n_certify, self.input_dim])
+                zeta2 = np.random.choice([-1, 1], size = (n_certify, 1))
+            zeta = np.concatenate([zeta1, zeta2], axis = 1)
+        else:
+            raise NotImplementedError
+
+        data = np.concatenate([X, y.reshape(-1, 1)], axis = 1)
+        # in practice, we always include the empirical data in the sampled support
+        zeta = np.concatenate([data, zeta])
+
+        kernel_width, kernel_gamma = self.medium_heuristic(zeta, zeta)
+        
+        # --------------------------------------------------------------------------------
+        # Step 3: setup objective function and constraints
+        # --------------------------------------------------------------------------------
+        # evaluate the f, K at the value of zetas
+        K = rbf_kernel(zeta, zeta)
+        #gamma = kernel_gamma)
+        f = a @ K
+        constr = []
+        for i in range((len(zeta))):
+            constr += [self.cvx_loss(theta, zeta[i]) <= f0 + f[i]]
+
+        obj = f0 + cp.sum(f[0:sample_size]) / sample_size + self.eta * cp.norm(a.T @ self.matrix_decomp(K))
+        opt = cp.Problem(cp.Minimize(obj), constr)
+    
+
+        opt.solve(solver = cp.MOSEK)
+        self.theta = theta.value #, obj.value, a.value, f0.value, kernel_gamma, zeta
+
+        

dro-0.0.1/dro/src/Sinkhorn_DRO.py

@@ -0,0 +1,277 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import TensorDataset, DataLoader
+from torch.autograd import Variable
+import math
+from sklearn.metrics import f1_score
+
+class LinearModel(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super(LinearModel, self).__init__()
+        self.linear = nn.Linear(input_dim, output_dim, bias=True)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+def to_tensor(x):
+    if type(x) == np.ndarray:
+        return torch.from_numpy(x).float()
+    elif type(x) == torch.Tensor:
+        return x
+    else:
+        print("Type error. Input should be either numpy array or torch tensor")
+
+def SDRO_eval(theta, Lambda, Reg, X, Y):
+    n,d = X.shape
+
+    ratio_1 = Lambda / (Lambda - 2 * np.linalg.norm(theta)**2)
+    residual = np.mean((X@theta - Y)**2,0)
+    obj_1 = ratio_1 * residual
+    obj_2 = Lambda*Reg/2 * np.linalg.slogdet(np.eye(d) - theta@theta.T*2/Lambda)[1]
+    return obj_1[0] - obj_2, np.sqrt(residual[0])
+
+
+
+
+class Sinkhorn_DRO_Linear:
+    def __init__(self, input_dim, reg_=1, lambda_=1, output_dim=1,
+                    maxiter = 50, learning_rate = 1e-2, K_sample_max=5, is_regression=0):
+        print(input_dim)
+        self.model = LinearModel(input_dim, output_dim)
+        self.lambda_ = lambda_
+        self.reg_ = reg_
+        self.maxiter_ = maxiter
+        self.learning_rate = learning_rate
+        self.K_sample_max = K_sample_max
+        self.is_regression = is_regression
+    
+    def update(self, config={}):
+        if "reg" in config.keys():
+            self.reg_ = config["reg"]
+        if "lambda" in config.keys():
+            self.lambda_ = config["lambda"]
+        if "k" in config.keys():
+            self.K_sample_max = config["k"]
+    
+    def predict(self, X):
+        X = torch.tensor(X).float()
+        self.model.cpu()
+        if self.is_regression:
+            pred = self.model(X)
+        else:
+            pred = self.model(X)
+            pred[pred > 0.5] = 1
+            pred[pred <= 0.5] = 0
+
+        return pred.detach().numpy()
+    
+    def score(self, X, y):
+        if self.is_regression:
+            return np.mean(self.predict(X).reshape(-1)-y.reshape(-1))
+        else:
+            acc = np.mean([self.predict(X).flatten() == y.flatten()])
+            f1 = f1_score(y, self.predict(X), average='macro')
+            return acc, f1 
+        
+    def fit(self, X, y, optimization_type="SG"):
+        X_tensor = torch.tensor(X, dtype=torch.float32)  
+        Y_tensor = torch.tensor(y, dtype=torch.float32)  
+        dataset = TensorDataset(X_tensor, Y_tensor)
+        dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
+
+        if optimization_type == 'SG':
+            self.SDRO_SG_solver(dataloader)
+        elif optimization_type == 'MLMC':
+            self.SDRO_MLMC_solver(dataloader)
+        elif optimization_type == 'RTMLMC':
+            self.SDRO_RTMLMC_solver(dataloader)
+        else:
+            raise ImplementationError("This optimization method is not implemented! Please choose one in \{SG, MLMC, RTMLMC\}")
+        pass
+
+        theta = self.model.linear.weight.cpu().detach().tolist()
+        bias = self.model.linear.bias.cpu().detach().tolist()
+
+        params = {}
+        params['theta'] = theta 
+        params['bias'] = bias 
+        return params
+
+
+    def SDRO_SG_solver(self, dataloader, device=torch.device("cuda:7")):
+        """
+        2-SDRO Approach with SG estimator for Regression Problem
+        #   Input:
+        # Feature: N samples of R^d [dim: N*d]
+        #  Target: labels of N samples [dim: N*1]
+        #   theta: initial guess for optimization
+        #  Lambda: Lagrangian multiplier
+        #     Reg: bandwidth
+        #  Output:
+        #   theta: optimized decision
+        """
+        iter = 0
+        Lambda_Reg = self.lambda_ * self.reg_
+        self.model.to(device)
+        optimizer_theta = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
+
+        
+        for epoch in range(self.maxiter_):
+            for _, (data, target) in enumerate(dataloader):
+                iter = iter + 1
+
+                # generate stochastic samples
+                N, d         = data.shape
+                data, target = Variable(data), Variable(target)
+                m            = int(2**self.K_sample_max)
+
+                optimizer_theta.zero_grad()
+                data_noise     = torch.randn([m, N, d]) * np.sqrt(self.reg_) + data.reshape([1,N,d])
+                data_noise_vec = data_noise.reshape([-1,d])
+                target_noise   = target.repeat(m,1).reshape(-1,1).to(device)
+
+                haty  = self.model(data_noise_vec.to(device))
+
+                obj_vec  = (haty - target_noise) ** 2
+                obj_mat  = obj_vec.reshape([m, N])
+                Residual = obj_mat / Lambda_Reg
+
+                Loss_SDRO = (torch.logsumexp(Residual, dim=0, keepdim=True)-math.log(m)) * Lambda_Reg
+                Loss_SDRO_avg = torch.mean(Loss_SDRO)
+
+                Loss_SDRO_avg.backward()
+                optimizer_theta.step()
+
+                if iter % 10000 == 0:
+                    print(f"Iter {iter} {Loss_SDRO_avg.data}")
+
+
+    def SDRO_MLMC_solver(self, dataloader):
+        iter       = 0
+        Lambda_Reg = self.lambda_ * self.reg_
+        
+        optimizer_theta = torch.optim.SGD(self.model.parameters(), lr=self.learning_rate)
+        N_ell_hist      = np.int_(2**(np.arange(self.K_sample_max) + 1))
+
+        
+        for epoch in range(self.maxiter_):
+            for _, (data, target) in enumerate(dataloader):
+                iter = iter + 1
+                # generate stochastic samples
+                N, d         = data.shape
+                data, target = Variable(data), Variable(target)
+                optimizer_theta.zero_grad()
+
+                m_total = 0
+                for K_sample in np.arange(self.K_sample_max):
+                    m          = int(2**K_sample)
+                    N_ell      = N_ell_hist[-K_sample-1]
+                    data_ell   = data[:N_ell, :]
+                    N_ell, d   = data_ell.shape
+                    target_ell = target[:N_ell]
+
+                    data_noise     = torch.randn([m, N_ell, d]) * np.sqrt(self.reg_) + data_ell.reshape([1,N_ell,d])
+                    m_total       += m * N_ell
+                    data_noise_vec = data_noise.reshape([-1,d])
+                    target_noise   = target_ell.repeat(m,1)
+
+                    haty     = self.model(data_noise_vec)
+                    obj_vec  = (haty - target_noise) ** 2
+                    obj_mat  = obj_vec.reshape([m, N_ell])
+                    Residual = obj_mat / Lambda_Reg
+
+                    if K_sample == 0:
+                        Loss_SDRO = torch.log(torch.mean(torch.exp(Residual), dim=0)) * Lambda_Reg
+                        Loss_SDRO_avg_K_sample = torch.mean(Loss_SDRO)
+                        Loss_SDRO_avg_sum = Loss_SDRO_avg_K_sample
+                    else:
+                        m1 = int(m/2)
+                        Residual_half   = Residual[:m1,:]
+                        Residual_remain = Residual[m1:,:]
+                        Loss_SDRO_1     = torch.log(torch.mean(torch.exp(Residual), dim=0)) * Lambda_Reg
+                        Loss_SDRO_avg_1 = torch.mean(Loss_SDRO_1)
+                        Loss_SDRO_2     = (torch.log(torch.mean(torch.exp(Residual_half), dim=0)) + torch.log(torch.mean(torch.exp(Residual_remain), dim=0))) * Lambda_Reg
+                        Loss_SDRO_avg_2 = torch.mean(Loss_SDRO_2)
+                        Loss_SDRO_avg_K_sample   = Loss_SDRO_avg_1 - 0.5 * Loss_SDRO_avg_2                
+                        Loss_SDRO_avg_sum = Loss_SDRO_avg_sum + Loss_SDRO_avg_K_sample
+
+                Loss_SDRO_avg_sum.backward()
+                optimizer_theta.step()
+                
+                # if torch.linalg.norm(self.model.linear.weight.data) > 0.95*np.sqrt(Lambda/2):
+                #     self.model.linear.weight.data = self.model.linear.weight.data / torch.linalg.norm(self.model.linear.weight.data) * 0.95*np.sqrt(Lambda/2)
+        
+        
+
+    def SDRO_RTMLMC_solver(self, dataloader):
+        iter            = 0
+        Lambda_Reg      = self.lambda_ * self.reg_
+        optimizer_theta = torch.optim.SGD(self.model.parameters(), lr=self.learning_rate)
+
+        # sampling from truncated gemoetric distribution
+        elements      = np.arange(self.K_sample_max)
+        probabilities = (0.5) ** (elements)
+        probabilities = probabilities / np.sum(probabilities)
+
+        for epoch in range(self.maxiter_):
+            for _, (data, target) in enumerate(dataloader):
+                iter = iter + 1
+                # generate stochastic samples
+                N, d         = data.shape
+                data, target = Variable(data), Variable(target)
+
+                K_sample       = int(np.random.choice(list(elements), 1, list(probabilities)))
+                m              = int(2**K_sample)
+                optimizer_theta.zero_grad()
+                data_noise     = torch.randn([m, N, d]) * np.sqrt(Reg) + data.reshape([1,N,d])
+                data_noise_vec = data_noise.reshape([-1,d])
+                target_noise   = target.repeat(m,1)
+
+                haty       = data_noise_vec @ theta_torch.type(torch.float64)
+                obj_vec    = (haty - target_noise) ** 2
+                obj_mat    = obj_vec.reshape([m, N])
+                Residual   = obj_mat / Lambda_Reg
+
+                if m == 1:
+                    Loss_SDRO = torch.log(torch.mean(torch.exp(Residual), dim=0)) * Lambda_Reg
+                    Loss_SDRO_avg = torch.mean(Loss_SDRO)
+                else:
+                    m1 = int(2**(K_sample-1))
+                    Residual_half   = Residual[:m1,:]
+                    Residual_remain = Residual[m1:,:]
+                    Loss_SDRO_1     = torch.log(torch.mean(torch.exp(Residual), dim=0)) * Lambda_Reg
+                    Loss_SDRO_avg_1 = torch.mean(Loss_SDRO_1)
+                    Loss_SDRO_2     = (torch.log(torch.mean(torch.exp(Residual_half), dim=0)) + torch.log(torch.mean(torch.exp(Residual_remain), dim=0))) * Lambda_Reg
+                    Loss_SDRO_avg_2 = torch.mean(Loss_SDRO_2)
+                    Loss_SDRO_avg   = Loss_SDRO_avg_1 - 0.5 * Loss_SDRO_avg_2
+                
+                Loss_SDRO_avg = 1/probabilities[K_sample]* Loss_SDRO_avg
+                Loss_SDRO_avg.backward()
+                optimizer_theta.step()
+        
+                # if torch.linalg.norm(self.model.linear.weight.data) > 0.95*np.sqrt(Lambda/2):
+                #     self.model.linear.weight.data = self.model.linear.weight.data / torch.linalg.norm(self.model.linear.weight.data) * 0.95*np.sqrt(Lambda/2)
+        
+
+
+
+if __name__ == "__main__":
+    from sklearn.datasets import make_regression
+    from  sklearn.linear_model import LinearRegression
+    import numpy as np
+
+    sample_size = 1000
+    feature_size = 10
+    X, y = make_regression(n_samples = sample_size, n_features = feature_size, noise = 1, random_state = 42)
+    
+    method = Sinkhorn_DRO_Linear(0.1, 1000.0, maxiter=2000, input_dim=feature_size)
+    method.fit(X,y)
+    print(method.model.linear.weight)
+
+    model = LinearRegression(fit_intercept=True)
+    model.fit(X, y)
+    print(model.coef_)