def test(): torch.manual_seed(1) torch.cuda.manual_seed_all(1) model = cnn_MNIST() checkpoint = torch.load("../examples/vision/pretrain/mnist_cnn_small.pth", map_location="cpu") model.load_state_dict(checkpoint) N = 2 n_classes = 10 image = torch.randn(N, 1, 28, 28) image = image.to(torch.float32) / 255.0 model = BoundedModule(model, torch.empty_like(image), device="cpu") eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) pred = model(image) lb, ub = model.compute_bounds() assert lb.shape == ub.shape == torch.Size((2, 10)) path = 'data/constant_test_data' if args.gen_ref: torch.save((lb, ub), path) else: lb_ref, ub_ref = torch.load(path) print(lb) print(lb_ref) assert torch.allclose(lb, lb_ref) assert torch.allclose(ub, ub_ref)
def test(self): model_oris = [ models.model_resnet(width=1, mult=2), models.ResNet18(in_planes=2) ] self.result = [] for model_ori in model_oris: conv_mode = 'patches' # conv_mode can be set as 'matrix' or 'patches' normalize = torchvision.transforms.Normalize( mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) test_data = torchvision.datasets.CIFAR10( "./data", train=False, download=True, transform=torchvision.transforms.Compose( [torchvision.transforms.ToTensor(), normalize])) N = 1 n_classes = 10 image = torch.Tensor(test_data.data[:N]).reshape(N, 3, 32, 32) image = image.to(torch.float32) / 255.0 model = BoundedModule(model_ori, image, bound_opts={"conv_mode": conv_mode}) ptb = PerturbationLpNorm(norm=np.inf, eps=0.03) image = BoundedTensor(image, ptb) pred = model(image) lb, ub = model.compute_bounds(IBP=False, C=None, method='backward') self.result += [lb, ub] self.check()
def compute_and_compare_bounds(self, eps, norm, IBP, method): input_data = torch.randn((N, 256)) model = BoundedModule(self.original_model, torch.empty_like(input_data)) ptb = PerturbationLpNorm(norm=norm, eps=eps) ptb_data = BoundedTensor(input_data, ptb) pred = model(ptb_data) label = torch.argmax(pred, dim=1).cpu().detach().numpy() # Compute bounds. lb, ub = model.compute_bounds(IBP=IBP, method=method) # Compute dual norm. if norm == 1: q = np.inf elif norm == np.inf: q = 1.0 else: q = 1.0 / (1.0 - (1.0 / norm)) # Compute reference manually. weight, bias = list(model.parameters()) norm = weight.norm(p=q, dim=1) expected_pred = input_data.matmul(weight.t()) + bias expected_ub = eps * norm + expected_pred expected_lb = -eps * norm + expected_pred # Check equivalence. self.assertEqual(expected_pred, pred) self.assertEqual(expected_ub, ub) self.assertEqual(expected_lb, lb)
def test(): torch.manual_seed(1) torch.cuda.manual_seed_all(1) models = [2, 3] paddings = [1, 2] strides = [1, 3] N = 2 n_classes = 10 image = torch.randn(N, 1, 28, 28) image = image.to(torch.float32) / 255.0 for layer_num in models: for padding in paddings: for stride in strides: # print(layer_num, padding, stride) try: model_ori = cnn_model(layer_num, padding, stride) except: continue model = BoundedModule(model_ori, torch.empty_like(image), device="cpu", bound_opts={"conv_mode": "patches"}) eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) pred = model(image) lb, ub = model.compute_bounds() model = BoundedModule(model_ori, torch.empty_like(image), device="cpu", bound_opts={"conv_mode": "matrix"}) pred = model(image) lb_ref, ub_ref = model.compute_bounds() assert lb.shape == ub.shape == torch.Size((N, n_classes)) assert torch.allclose(lb, lb_ref) assert torch.allclose(ub, ub_ref)
def test(): net = ResNet18() N = 2 n_classes = 10 x = torch.randn(N, 3, 32, 32) y = net(x) device = 'cpu' if device == 'cuda': x = x.cuda() y = y.cuda() model = BoundedModule(net, torch.empty_like(x), bound_opts={"conv_mode": "patches"}, device=device) print("Model structure: \n", str(net)) eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(x, ptb) pred = model(image) lb, ub = model.compute_bounds() model = BoundedModule(net, torch.empty_like(x), bound_opts={"conv_mode": "matrix"}, device=device) eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(x, ptb) pred = model(image) lb_ref, ub_ref = model.compute_bounds() # assert lb.shape == ub.shape == torch.Size((N, n_classes)) print((lb - lb_ref).sum(), (ub - ub_ref).sum()) assert torch.allclose(lb, lb_ref) assert torch.allclose(ub, ub_ref)
def compute_and_compare_bounds(self, eps, norm, IBP, method): input_data = torch.randn((N, 1, input_dim, input_dim)) model = BoundedModule(self.original_model, torch.empty_like(input_data)) ptb = PerturbationLpNorm(norm=norm, eps=eps) ptb_data = BoundedTensor(input_data, ptb) pred = model(ptb_data) label = torch.argmax(pred, dim=1).cpu().detach().numpy() # Compute bounds. lb, ub = model.compute_bounds(IBP=IBP, method=method) # Compute reference. conv_weight, conv_bias = list(model.parameters()) conv_bias = conv_bias.view(1, out_channel, 1, 1) matrix_eye = torch.eye(input_dim * input_dim).view( input_dim * input_dim, 1, input_dim, input_dim) # Obtain equivalent weight and bias for convolution. weight = self.original_model.conv( matrix_eye ) - conv_bias # Output is (batch, channel, weight, height). weight = weight.view( input_dim * input_dim, -1) # Dimension is (flattened_input, flattened_output). bias = conv_bias.repeat(1, 1, input_dim // 2, input_dim // 2).view(-1) flattend_data = input_data.view(N, -1) # Compute dual norm. if norm == 1: q = np.inf elif norm == np.inf: q = 1.0 else: q = 1.0 / (1.0 - (1.0 / norm)) # Manually compute bounds. norm = weight.t().norm(p=q, dim=1) expected_pred = flattend_data.matmul(weight) + bias expected_ub = eps * norm + expected_pred expected_lb = -eps * norm + expected_pred # Check equivalence. if method == 'backward' or method == 'forward': self.assertEqual(expected_pred, pred) self.assertEqual(expected_ub, ub) self.assertEqual(expected_lb, lb)
def test(self): model = cnn_MNIST() checkpoint = torch.load( "../examples/vision/pretrain/mnist_cnn_small.pth", map_location="cpu") model.load_state_dict(checkpoint) N = 2 n_classes = 10 image = torch.randn(N, 1, 28, 28) image = image.to(torch.float32) / 255.0 model = BoundedModule(model, torch.empty_like(image), device="cpu") eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) pred = model(image) lb, ub = model.compute_bounds() assert lb.shape == ub.shape == torch.Size((2, 10)) self.result = (lb, ub) self.check()
image, bound_opts={"conv_mode": conv_mode}, device=device) ## Step 4: Compute bounds using LiRPA given a perturbation eps = 0.03 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) # Get model prediction as usual pred = model(image) # Compute bounds torch.cuda.empty_cache() print('Using {} mode to compute convolution.'.format(conv_mode)) lb, ub = model.compute_bounds(IBP=False, C=None, method='backward') ## Step 5: Final output # pred = pred.detach().cpu().numpy() lb = lb.detach().cpu().numpy() ub = ub.detach().cpu().numpy() for i in range(N): # print("Image {} top-1 prediction {}".format(i, label[i])) for j in range(n_classes): print("f_{j}(x_0): {l:8.3f} <= f_{j}(x_0+delta) <= {u:8.3f}".format( j=j, l=lb[i][j], u=ub[i][j])) print() # Print the GPU memory usage print('Memory usage in "{}" mode:'.format(conv_mode)) print(torch.cuda.memory_summary())
if batch_idx < 2: ## Step 3: wrap model with auto_LiRPA # The second parameter is for constructing the trace of the computational graph, and its content is not important. model = BoundedModule(model, inputs, device="cuda") ## Step 4: Compute bounds using LiRPA given a perturbation eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(inputs, ptb) # Get model prediction as usual pred = model(image) label = torch.argmax(pred, dim=1).cpu().numpy() # Compute bounds lb, ub = model.compute_bounds() ## Step 5: Final output pred = pred.detach().cpu().numpy() lb = lb.detach().cpu().numpy() ub = ub.detach().cpu().numpy() for i in range(N): print("Image {} top-1 prediction {}".format(i, label[i])) for j in range(n_classes): print( "f_{j}(x_0) = {fx0:8.3f}, {l:8.3f} <= f_{j}(x_0+delta) <= {u:8.3f}" .format(j=j, fx0=pred[i][j], l=lb[i][j], u=ub[i][j])) print()
class LiRPAConvNet: def __init__(self, model_ori, pred, test, solve_slope=False, device='cuda', simplify=True, in_size=(1, 3, 32, 32)): """ convert pytorch model to auto_LiRPA module """ layers = list(model_ori.children()) if simplify: added_prop_layers = add_single_prop(layers, pred, test) self.layers = added_prop_layers else: self.layers = layers net = nn.Sequential(*self.layers) self.solve_slope = solve_slope if solve_slope: self.net = BoundedModule(net, torch.rand(in_size), bound_opts={ 'relu': 'random_evaluation', 'conv_mode': 'patches' }, device=device) else: self.net = BoundedModule(net, torch.rand(in_size), bound_opts={'relu': 'same-slope'}, device=device) self.net.eval() def get_lower_bound(self, pre_lbs, pre_ubs, decision, slopes=None, history=[], decision_thresh=0, layer_set_bound=True, beta=True): """ # (in) pre_lbs: layers list -> tensor(batch, layer shape) # (in) relu_mask: relu layers list -> tensor(batch, relu layer shape (view-1)) # (in) slope: relu layers list -> tensor(batch, relu layer shape) # (out) lower_bounds: batch list -> layers list -> tensor(layer shape) # (out) masks_ret: batch list -> relu layers list -> tensor(relu layer shape) # (out) slope: batch list -> relu layers list -> tensor(relu layer shape) """ start = time.time() lower_bounds, upper_bounds, masks_ret, slopes = self.update_bounds_parallel( pre_lbs, pre_ubs, decision, slopes, beta=beta, early_stop=False, opt_choice="adam", iteration=20, history=history, decision_thresh=decision_thresh, layer_set_bound=layer_set_bound) end = time.time() print('batch time: ', end - start) return [i[-1] for i in upper_bounds ], [i[-1] for i in lower_bounds ], None, masks_ret, lower_bounds, upper_bounds, slopes def get_relu(self, model, idx): # find the i-th ReLU layer i = 0 for layer in model.children(): if isinstance(layer, BoundRelu): i += 1 if i == idx: return layer def get_candidate(self, model, lb, ub): # get the intermediate bounds in the current model and build self.name_dict which contains the important index # and model name pairs if self.input_domain.ndim == 2: lower_bounds = [self.input_domain[:, 0].squeeze(-1)] upper_bounds = [self.input_domain[:, 1].squeeze(-1)] else: lower_bounds = [self.input_domain[:, :, :, 0].squeeze(-1)] upper_bounds = [self.input_domain[:, :, :, 1].squeeze(-1)] self.pre_relu_indices = [] idx, i, model_i = 0, 0, 0 # build a name_dict to map layer idx in self.layers to BoundedModule self.name_dict = {0: model.root_name[0]} model_names = list(model._modules) for layer in self.layers: if isinstance(layer, nn.ReLU): i += 1 this_relu = self.get_relu(model, i) lower_bounds[-1] = this_relu.inputs[0].lower.squeeze().detach() upper_bounds[-1] = this_relu.inputs[0].upper.squeeze().detach() lower_bounds.append(F.relu(lower_bounds[-1]).detach()) upper_bounds.append(F.relu(upper_bounds[-1]).detach()) self.pre_relu_indices.append(idx) self.name_dict[idx + 1] = model_names[model_i] model_i += 1 elif isinstance(layer, Flatten): lower_bounds.append(lower_bounds[-1].reshape(-1).detach()) upper_bounds.append(upper_bounds[-1].reshape(-1).detach()) self.name_dict[idx + 1] = model_names[model_i] model_i += 8 # Flatten is split to 8 ops in BoundedModule elif isinstance(layer, ZeroPad2d): lower_bounds.append(F.pad(lower_bounds[-1], layer.padding)) upper_bounds.append(F.pad(upper_bounds[-1], layer.padding)) self.name_dict[idx + 1] = model_names[model_i] model_i += 24 else: self.name_dict[idx + 1] = model_names[model_i] lower_bounds.append([]) upper_bounds.append([]) model_i += 1 idx += 1 # Also add the bounds on the final thing lower_bounds[-1] = (lb.view(-1).detach()) upper_bounds[-1] = (ub.view(-1).detach()) return lower_bounds, upper_bounds, self.pre_relu_indices def get_candidate_parallel(self, model, lb, ub, batch): # get the intermediate bounds in the current model lower_bounds = [ self.input_domain[:, :, :, 0].squeeze(-1).repeat(batch, 1, 1, 1) ] upper_bounds = [ self.input_domain[:, :, :, 1].squeeze(-1).repeat(batch, 1, 1, 1) ] idx, i, = 0, 0 for layer in self.layers: if isinstance(layer, nn.ReLU): i += 1 this_relu = self.get_relu(model, i) lower_bounds[-1] = this_relu.inputs[0].lower.detach() upper_bounds[-1] = this_relu.inputs[0].upper.detach() lower_bounds.append(F.relu(lower_bounds[-1]).detach( )) # TODO we actually do not need the bounds after ReLU upper_bounds.append(F.relu(upper_bounds[-1]).detach()) elif isinstance(layer, Flatten): lower_bounds.append(lower_bounds[-1].reshape(batch, -1).detach()) upper_bounds.append(upper_bounds[-1].reshape(batch, -1).detach()) elif isinstance(layer, nn.ZeroPad2d): lower_bounds.append( F.pad(lower_bounds[-1], layer.padding).detach()) upper_bounds.append( F.pad(upper_bounds[-1], layer.padding).detach()) else: lower_bounds.append([]) upper_bounds.append([]) idx += 1 # Also add the bounds on the final thing lower_bounds[-1] = (lb.view(batch, -1).detach()) upper_bounds[-1] = (ub.view(batch, -1).detach()) return lower_bounds, upper_bounds def get_mask_parallel(self, model): # get the mask of status of ReLU, 0 means inactive neurons, -1 means unstable neurons, 1 means active neurons mask = [] idx, i, = 0, 0 for layer in self.layers: if isinstance(layer, nn.ReLU): i += 1 this_relu = self.get_relu(model, i) mask_tmp = torch.zeros_like(this_relu.inputs[0].lower) unstable = ((this_relu.inputs[0].lower < 0) & (this_relu.inputs[0].upper > 0)) mask_tmp[unstable] = -1 active = (this_relu.inputs[0].lower >= 0) mask_tmp[active] = 1 # otherwise 0, for inactive neurons mask.append(mask_tmp.reshape(mask_tmp.size(0), -1)) ret = [] for i in range(mask[0].size(0)): ret.append([j[i] for j in mask]) return ret def get_beta(self, model): b = [] bm = [] for m in model._modules.values(): if isinstance(m, BoundRelu): b.append(m.beta.clone().detach()) bm.append(m.beta_mask.clone().detach()) retb = [] retbm = [] for i in range(b[0].size(0)): retb.append([j[i] for j in b]) retbm.append([j[i] for j in bm]) return (retb, retbm) def get_slope(self, model): s = [] for m in model._modules.values(): if isinstance(m, BoundRelu): s.append(m.slope.transpose(0, 1).clone().detach()) ret = [] for i in range(s[0].size(0)): ret.append([j[i] for j in s]) return ret def set_slope(self, model, slope): idx = 0 for m in model._modules.values(): if isinstance(m, BoundRelu): # m.slope = slope[idx].repeat(2, *([1] * (slope[idx].ndim - 1))).requires_grad_(True) m.slope = slope[idx].repeat( 2, *([1] * (slope[idx].ndim - 1))).transpose( 0, 1).requires_grad_(True) idx += 1 def reset_beta(self, model, batch=0): if batch == 0: for m in model._modules.values(): if isinstance(m, BoundRelu): m.beta.data = m.beta.data * 0. m.beta_mask.data = m.beta_mask.data * 0. # print("beta[{}]".format(batch), m.beta.shape, m.beta_mask.shape) else: for m in model._modules.values(): if isinstance(m, BoundRelu): ndim = m.beta.data.ndim # m.beta.data=(m.beta.data[0:1]*0.).repeat(batch*2, *([1] * (ndim - 1))).requires_grad_(True) # m.beta_mask.data=(m.beta_mask.data[0:1]*0.).repeat(batch*2, *([1] * (ndim - 1))).requires_grad_(True) m.beta = torch.zeros(m.beta[:, 0:1].shape).repeat( 1, batch * 2, *([1] * (ndim - 2))).detach().to( m.beta.device).requires_grad_(True) m.beta_mask = torch.zeros(m.beta_mask[0:1].shape).repeat( batch * 2, *([1] * (ndim - 2))).detach().to( m.beta.device).requires_grad_(False) # print("beta[{}]".format(batch), m.beta.shape, m.beta_mask.shape) def update_bounds_parallel(self, pre_lb_all=None, pre_ub_all=None, decision=None, slopes=None, beta=True, early_stop=True, opt_choice="default", iteration=20, history=[], decision_thresh=0, layer_set_bound=True): # update optimize-CROWN bounds in a parallel way total_batch = len(decision) decision = np.array(decision) layers_need_change = np.unique(decision[:, 0]) layers_need_change.sort() # initial results with empty list ret_l = [[] for _ in range(len(decision) * 2)] ret_u = [[] for _ in range(len(decision) * 2)] masks = [[] for _ in range(len(decision) * 2)] ret_s = [[] for _ in range(len(decision) * 2)] pre_lb_all_cp = copy.deepcopy(pre_lb_all) pre_ub_all_cp = copy.deepcopy(pre_ub_all) for idx in layers_need_change: # iteratively change upper and lower bound from former to later layer tmp_d = np.argwhere(decision[:, 0] == idx) # .squeeze() # idx is the index of relu layers, change_idx is the index of all layers change_idx = self.pre_relu_indices[idx] batch = len(tmp_d) select_history = [ history[idx] for idx in tmp_d.squeeze().reshape(-1) ] if beta: # update beta mask, put it after reset_beta # reset beta according to the shape of batch self.reset_beta(self.net, batch) # print("select history", select_history) bound_relus = [] for m in self.net._modules.values(): if isinstance(m, BoundRelu): bound_relus.append(m) m.beta_mask.data = m.beta_mask.data.view(batch * 2, -1) for bi in range(batch): d = tmp_d[bi][0] # assign current decision to each point of a batch bound_relus[int(decision[d][0])].beta_mask.data[ bi, int(decision[d][1])] = 1 bound_relus[int(decision[d][0])].beta_mask.data[ bi + batch, int(decision[d][1])] = -1 # print("assign", bi, decision[d], 1, bound_relus[decision[d][0]].beta_mask.data[bi, decision[d][1]]) # print("assign", bi+batch, decision[d], -1, bound_relus[decision[d][0]].beta_mask.data[bi+batch, decision[d][1]]) # assign history decision according to select_history for (hid, hl), hc in select_history[bi]: bound_relus[hid].beta_mask.data[bi, hl] = int( (hc - 0.5) * 2) bound_relus[hid].beta_mask.data[bi + batch, hl] = int( (hc - 0.5) * 2) # print("assign", bi, [hid, hl], hc, bound_relus[hid].beta_mask.data[bi, hl]) # print("assign", bi+batch, [hid, hl], hc, bound_relus[hid].beta_mask.data[bi+batch, hl]) # sanity check: beta_mask should only be assigned for split nodes for m in bound_relus: m.beta_mask.data = m.beta_mask.data.view(m.beta[0].shape) slope_select = [i[tmp_d.squeeze()].clone() for i in slopes] pre_lb_all = [i[tmp_d.squeeze()].clone() for i in pre_lb_all_cp] pre_ub_all = [i[tmp_d.squeeze()].clone() for i in pre_ub_all_cp] if batch == 1: pre_lb_all = [i.clone().unsqueeze(0) for i in pre_lb_all] pre_ub_all = [i.clone().unsqueeze(0) for i in pre_ub_all] slope_select = [i.clone().unsqueeze(0) for i in slope_select] upper_bounds = [i.clone() for i in pre_ub_all[:change_idx + 1]] lower_bounds = [i.clone() for i in pre_lb_all[:change_idx + 1]] upper_bounds_cp = copy.deepcopy(upper_bounds) lower_bounds_cp = copy.deepcopy(lower_bounds) for i in range(batch): d = tmp_d[i][0] upper_bounds[change_idx].view(batch, -1)[i][decision[d][1]] = 0 lower_bounds[change_idx].view(batch, -1)[i][decision[d][1]] = 0 pre_lb_all = [torch.cat(2 * [i]) for i in pre_lb_all] pre_ub_all = [torch.cat(2 * [i]) for i in pre_ub_all] # merge the inactive and active splits together new_candidate = {} for i, (l, uc, lc, u) in enumerate( zip(lower_bounds, upper_bounds_cp, lower_bounds_cp, upper_bounds)): # we set lower = 0 in first half batch, and upper = 0 in second half batch new_candidate[self.name_dict[i]] = [ torch.cat((l, lc), dim=0), torch.cat((uc, u), dim=0) ] if not layer_set_bound: new_candidate_p = {} for i, (l, u) in enumerate(zip(pre_lb_all[:-2], pre_ub_all[:-2])): # we set lower = 0 in first half batch, and upper = 0 in second half batch new_candidate_p[self.name_dict[i]] = [l, u] # create new_x here since batch may change ptb = PerturbationLpNorm( norm=self.x.ptb.norm, eps=self.x.ptb.eps, x_L=self.x.ptb.x_L.repeat(batch * 2, 1, 1, 1), x_U=self.x.ptb.x_U.repeat(batch * 2, 1, 1, 1)) new_x = BoundedTensor(self.x.data.repeat(batch * 2, 1, 1, 1), ptb) self.net( new_x ) # batch may change, so we need to do forward to set some shapes here if len(slope_select) > 0: # set slope here again self.set_slope(self.net, slope_select) torch.cuda.empty_cache() if layer_set_bound: # we fix the intermediate bounds before change_idx-th layer by using normal CROWN if self.solve_slope and change_idx >= self.pre_relu_indices[-1]: # we split the ReLU at last layer, directly use Optimized CROWN self.net.set_bound_opts({ 'ob_start_idx': sum(change_idx <= x for x in self.pre_relu_indices), 'ob_beta': beta, 'ob_update_by_layer': layer_set_bound, 'ob_iteration': iteration }) lb, ub, = self.net.compute_bounds( x=(new_x, ), IBP=False, C=None, method='CROWN-Optimized', new_interval=new_candidate, return_A=False, bound_upper=False) else: # we split the ReLU before the last layer, calculate intermediate bounds by using normal CROWN self.net.set_relu_used_count( sum(change_idx <= x for x in self.pre_relu_indices)) with torch.no_grad(): lb, ub, = self.net.compute_bounds( x=(new_x, ), IBP=False, C=None, method='backward', new_interval=new_candidate, bound_upper=False, return_A=False) # we don't care about the upper bound of the last layer lower_bounds_new, upper_bounds_new = self.get_candidate_parallel( self.net, lb, lb + 99, batch * 2) if change_idx < self.pre_relu_indices[-1]: # check whether we have a better bounds before, and preset all intermediate bounds for i, (l, u) in enumerate( zip(lower_bounds_new[change_idx + 2:-1], upper_bounds_new[change_idx + 2:-1])): new_candidate[self.name_dict[i + change_idx + 2]] = [ torch.max(l, pre_lb_all[i + change_idx + 2]), torch.min(u, pre_ub_all[i + change_idx + 2]) ] if self.solve_slope: self.net.set_bound_opts({ 'ob_start_idx': sum(change_idx <= x for x in self.pre_relu_indices), 'ob_beta': beta, 'ob_update_by_layer': layer_set_bound, 'ob_iteration': iteration }) lb, ub, = self.net.compute_bounds( x=(new_x, ), IBP=False, C=None, method='CROWN-Optimized', new_interval=new_candidate, return_A=False, bound_upper=False) else: self.net.set_relu_used_count( sum(change_idx <= x for x in self.pre_relu_indices)) with torch.no_grad(): lb, ub, = self.net.compute_bounds( x=(new_x, ), IBP=False, C=None, method='backward', new_interval=new_candidate, bound_upper=False, return_A=False) else: # all intermediate bounds are re-calculate by optimized CROWN self.net.set_bound_opts({ 'ob_start_idx': 99, 'ob_beta': beta, 'ob_update_by_layer': layer_set_bound, 'ob_iteration': iteration }) lb, ub, = self.net.compute_bounds(x=(new_x, ), IBP=False, C=None, method='CROWN-Optimized', new_interval=new_candidate_p, return_A=False, bound_upper=False) # print('best results of parent nodes', pre_lb_all[-1].repeat(2, 1)) # print('finally, after optimization:', lower_bounds_new[-1]) # primal = self.get_primals(A_dict, return_x=True) lower_bounds_new, upper_bounds_new = self.get_candidate_parallel( self.net, lb, lb + 99, batch * 2) lower_bounds_new[-1] = torch.max(lower_bounds_new[-1], pre_lb_all[-1]) upper_bounds_new[-1] = torch.min(upper_bounds_new[-1], pre_ub_all[-1]) mask = self.get_mask_parallel(self.net) if len(slope_select) > 0: slope = self.get_slope(self.net) # reshape the results for i in range(len(tmp_d)): ret_l[int(tmp_d[i])] = [j[i] for j in lower_bounds_new] ret_l[int(tmp_d[i] + total_batch)] = [ j[i + batch] for j in lower_bounds_new ] ret_u[int(tmp_d[i])] = [j[i] for j in upper_bounds_new] ret_u[int(tmp_d[i] + total_batch)] = [ j[i + batch] for j in upper_bounds_new ] masks[int(tmp_d[i])] = mask[i] masks[int(tmp_d[i] + total_batch)] = mask[i + batch] if len(slope_select) > 0: ret_s[int(tmp_d[i])] = slope[i] ret_s[int(tmp_d[i] + total_batch)] = slope[i + batch] return ret_l, ret_u, masks, ret_s def fake_forward(self, x): for layer in self.layers: if type(layer) is nn.Linear: x = F.linear(x, layer.weight, layer.bias) elif type(layer) is nn.Conv2d: x = F.conv2d(x, layer.weight, layer.bias, layer.stride, layer.padding, layer.dilation, layer.groups) elif type(layer) == nn.ReLU: x = F.relu(x) elif type(layer) == Flatten: x = x.reshape(x.shape[0], -1) elif type(layer) == nn.ZeroPad2d: x = F.pad(x, layer.padding) else: print(type(layer)) raise NotImplementedError return x def get_primals(self, A, return_x=False): # get primal input by using A matrix input_A_lower = A[self.layer_names[-1]][self.net.input_name[0]][0] batch = input_A_lower.shape[1] l = self.input_domain[:, :, :, 0].repeat(batch, 1, 1, 1) u = self.input_domain[:, :, :, 1].repeat(batch, 1, 1, 1) diff = 0.5 * (l - u) # already flip the sign by using lower - upper net_input = diff * torch.sign(input_A_lower.squeeze(0)) + self.x if return_x: return net_input primals = [net_input] for layer in self.layers: if type(layer) is nn.Linear: pre = primals[-1] primals.append(F.linear(pre, layer.weight, layer.bias)) elif type(layer) is nn.Conv2d: pre = primals[-1] primals.append( F.conv2d(pre, layer.weight, layer.bias, layer.stride, layer.padding, layer.dilation, layer.groups)) elif type(layer) == nn.ReLU: primals.append(F.relu(primals[-1])) elif type(layer) == Flatten: primals.append(primals[-1].reshape(primals[-1].shape[0], -1)) else: print(type(layer)) raise NotImplementedError # primals = primals[1:] primals = [i.detach().clone() for i in primals] # print('primals', primals[-1]) return net_input, primals def get_relu_mask(self): relu_mask = [] relu_idx = 0 for layer in self.layers: if type(layer) == nn.ReLU: relu_idx += 1 this_relu = self.get_relu(self.net, relu_idx) new_layer_mask = [] ratios_all = this_relu.d.squeeze(0) for slope in ratios_all.flatten(): if slope.item() == 1.0: new_layer_mask.append(1) elif slope.item() == 0.0: new_layer_mask.append(0) else: new_layer_mask.append(-1) relu_mask.append( torch.tensor(new_layer_mask).to(self.x.device)) return relu_mask def build_the_model(self, input_domain, x, no_lp=True, decision_thresh=0): self.x = x self.input_domain = input_domain slope_opt = None # first get CROWN bounds if self.solve_slope: self.net.init_slope(self.x) self.net.set_bound_opts({ 'ob_iteration': 100, 'ob_beta': False, 'ob_alpha': True, 'ob_opt_choice': "adam", 'ob_decision_thresh': decision_thresh, 'ob_early_stop': False, 'ob_log': False, 'ob_start_idx': 99, 'ob_keep_best': True, 'ob_update_by_layer': True, 'ob_lr': 0.1 }) lb, ub, A_dict = self.net.compute_bounds(x=(x, ), IBP=False, C=None, method='CROWN-Optimized', return_A=True, bound_upper=False) slope_opt = self.get_slope( self.net)[0] # initial with one node only else: with torch.no_grad(): lb, ub, A_dict = self.net.compute_bounds(x=(x, ), IBP=False, C=None, method='backward', return_A=True) # build a complete A_dict self.layer_names = list(A_dict[list(A_dict.keys())[-1]].keys())[2:] self.layer_names.sort() # update bounds print('initial CROWN bounds:', lb, ub) primals, mini_inp = None, None # mini_inp, primals = self.get_primals(self.A_dict) lb, ub, pre_relu_indices = self.get_candidate( self.net, lb, lb + 99) # primals are better upper bounds duals = None return ub[-1], lb[-1], mini_inp, duals, primals, self.get_relu_mask( ), lb, ub, pre_relu_indices, slope_opt
class mynet(nn.Module): def __init__(self): super(mynet, self).__init__() self.output = nn.sequential(nn.Linear(5, 10), nn.ReLU(), nn.Linear(10, 3)) def forward(self, input): return self.features(input) raw_model = mynet() bound_model = BoundedModule(raw_model, input_vec) num_actions = 3 batchsize = 5 label = torch.tensor([0, 2, 1, 1, 0]) bnd_state = BoundedTensor(input_vec, PerturbationLpNorm(norm=np.inf, eps=0.1)) c = torch.eye(3).type_as(input_vec)[label].unsqueeze(1) - torch.eye(3).type_as( input_vec).unsqueeze(0) I = (~(label.data.unsqueeze(1) == torch.arange(3).type_as( label.data).unsqueeze(0))) c = (c[I].view(input_vec.size(0), 2, 3)) pred = bound_model(input_vec) basic_bound, _ = bound_model.compute_bounds(IBP=False, method='backward') advance_bound, _ = bound_model.compute_bounds(C=c, IBP=False, method='backward') print(basic_bound.detach().numpy()) print(advance_bound.detach().numpy())
bound_opts={"conv_mode": conv_mode}, device="cuda") ## Step 4: Compute bounds using LiRPA given a perturbation eps = 0.03 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) # Get model prediction as usual pred = model(image) # Compute bounds torch.cuda.empty_cache() print('Using {} mode to compute convolution.'.format(conv_mode)) lb, ub = model.compute_bounds(x=(image, ), IBP=False, C=None, method='backward') ## Step 5: Final output # pred = pred.detach().cpu().numpy() lb = lb.detach().cpu().numpy() ub = ub.detach().cpu().numpy() for i in range(N): # print("Image {} top-1 prediction {}".format(i, label[i])) for j in range(n_classes): print("f_{j}(x_0): {l:8.3f} <= f_{j}(x_0+delta) <= {u:8.3f}".format( j=j, l=lb[i][j], u=ub[i][j])) print() # Print the GPU memory usage print('Memory usage in "{}" mode:'.format(conv_mode))
image = test_data.data[:N].view(N, 1, 28, 28).cuda() # Convert to float image = image.to(torch.float32) / 255.0 ## Step 3: wrap model with auto_LiRPA # The second parameter is for constructing the trace of the computational graph, and its content is not important. model = BoundedModule(model, torch.empty_like(image), device="cuda") ## Step 4: Compute bounds using LiRPA given a perturbation eps = 0.3 norm = np.inf # ptb = PerturbationL0Norm(eps=eps) ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) # Get model prediction as usual pred = model(image) # label = torch.argmax(pred, dim=1).cpu().numpy() # Compute bounds lb, ub = model.compute_bounds(IBP=False) ## Step 5: Final output # pred = pred.detach().cpu().numpy() lb = lb.detach().cpu().numpy() ub = ub.detach().cpu().numpy() for i in range(N): # print("Image {} top-1 prediction {}".format(i, label[i])) for j in range(n_classes): print("f_{j}(x_0): {l:8.3f} <= f_{j}(x_0+delta) <= {u:8.3f}".format( j=j, l=lb[i][j], u=ub[i][j])) print()
class Trainer(): ''' This is a class representing a Policy Gradient trainer, which trains both a deep Policy network and a deep Value network. Exposes functions: - advantage_and_return - multi_actor_step - reset_envs - run_trajectories - train_step Trainer also handles all logging, which is done via the "cox" library ''' def __init__(self, policy_net_class, value_net_class, params, store, advanced_logging=True, log_every=5): ''' Initializes a new Trainer class. Inputs; - policy, the class of policy network to use (inheriting from nn.Module) - val, the class of value network to use (inheriting from nn.Module) - step, a reference to a function to use for the policy step (see steps.py) - params, an dictionary with all of the required hyperparameters ''' # Parameter Loading self.params = Parameters(params) # Whether or not the value network uses the current timestep time_in_state = self.VALUE_CALC == "time" # Whether to use GPU (as opposed to CPU) if not self.CPU: torch.set_default_tensor_type("torch.cuda.FloatTensor") # Environment Loading def env_constructor(): # Whether or not we should add the time to the state horizon_to_feed = self.T if time_in_state else None return Env(self.GAME, norm_states=self.NORM_STATES, norm_rewards=self.NORM_REWARDS, params=self.params, add_t_with_horizon=horizon_to_feed, clip_obs=self.CLIP_OBSERVATIONS, clip_rew=self.CLIP_REWARDS, show_env=self.SHOW_ENV, save_frames=self.SAVE_FRAMES, save_frames_path=self.SAVE_FRAMES_PATH) self.envs = [env_constructor() for _ in range(self.NUM_ACTORS)] self.params.AGENT_TYPE = "discrete" if self.envs[0].is_discrete else "continuous" self.params.NUM_ACTIONS = self.envs[0].num_actions self.params.NUM_FEATURES = self.envs[0].num_features self.policy_step = step_with_mode(self.MODE) self.params.MAX_KL_INCREMENT = (self.params.MAX_KL_FINAL - self.params.MAX_KL) / self.params.TRAIN_STEPS self.advanced_logging = advanced_logging self.n_steps = 0 self.log_every = log_every # Instantiation self.policy_model = policy_net_class(self.NUM_FEATURES, self.NUM_ACTIONS, self.INITIALIZATION, time_in_state=time_in_state, activation=self.policy_activation) # Instantiate convex relaxation model when mode is 'robust_ppo' if self.MODE == 'robust_ppo': self.create_relaxed_model(time_in_state) opts_ok = (self.PPO_LR == -1 or self.PPO_LR_ADAM == -1) assert opts_ok, "One of ppo_lr and ppo_lr_adam must be -1 (off)." # Whether we should use Adam or simple GD to optimize the policy parameters if self.PPO_LR_ADAM != -1: kwargs = { 'lr':self.PPO_LR_ADAM, } if self.params.ADAM_EPS > 0: kwargs['eps'] = self.ADAM_EPS self.params.POLICY_ADAM = optim.Adam(self.policy_model.parameters(), **kwargs) else: self.params.POLICY_ADAM = optim.SGD(self.policy_model.parameters(), lr=self.PPO_LR) # If using a time dependent value function, add one extra feature # for the time ratio t/T if time_in_state: self.params.NUM_FEATURES = self.NUM_FEATURES + 1 # Value function optimization self.val_model = value_net_class(self.NUM_FEATURES, self.INITIALIZATION) self.val_opt = optim.Adam(self.val_model.parameters(), lr=self.VAL_LR, eps=1e-5) assert self.policy_model.discrete == (self.AGENT_TYPE == "discrete") # Learning rate annealing # From OpenAI hyperparametrs: # Set adam learning rate to 3e-4 * alpha, where alpha decays from 1 to 0 over training if self.ANNEAL_LR: lam = lambda f: 1-f/self.TRAIN_STEPS ps = optim.lr_scheduler.LambdaLR(self.POLICY_ADAM, lr_lambda=lam) vs = optim.lr_scheduler.LambdaLR(self.val_opt, lr_lambda=lam) self.params.POLICY_SCHEDULER = ps self.params.VALUE_SCHEDULER = vs if store is not None: self.setup_stores(store) else: print("Not saving results to cox store.") def create_relaxed_model(self, time_in_state=False): # Create state perturbation model for robust PPO training. if isinstance(self.policy_model, CtsPolicy): from .convex_relaxation import RelaxedCtsPolicyForState relaxed_policy_model = RelaxedCtsPolicyForState( self.NUM_FEATURES, self.NUM_ACTIONS, time_in_state=time_in_state, activation=self.policy_activation, policy_model=self.policy_model) dummy_input1 = torch.randn(1, self.NUM_FEATURES) inputs = (dummy_input1, ) self.relaxed_policy_model = BoundedModule(relaxed_policy_model, inputs) self.robust_eps_scheduler = LinearScheduler(self.params.ROBUST_PPO_EPS, self.params.ROBUST_PPO_EPS_SCHEDULER_OPTS) if self.params.ROBUST_PPO_BETA_SCHEDULER_OPTS == "same": self.robust_beta_scheduler = LinearScheduler(self.params.ROBUST_PPO_BETA, self.params.ROBUST_PPO_EPS_SCHEDULER_OPTS) else: self.robust_beta_scheduler = LinearScheduler(self.params.ROBUST_PPO_BETA, self.params.ROBUST_PPO_BETA_SCHEDULER_OPTS) else: raise NotImplementedError """Initialize sarsa training.""" def setup_sarsa(self, lr_schedule, eps_scheduler, beta_scheduler): # Create the Sarsa model, with S and A as the input. self.sarsa_model = ValueDenseNet(self.NUM_FEATURES + self.NUM_ACTIONS, self.INITIALIZATION) self.sarsa_opt = optim.Adam(self.sarsa_model.parameters(), lr=self.VAL_LR, eps=1e-5) self.sarsa_scheduler = optim.lr_scheduler.LambdaLR(self.sarsa_opt, lr_schedule) self.sarsa_eps_scheduler = eps_scheduler self.sarsa_beta_scheduler = beta_scheduler # Convert model with relaxation wrapper. dummy_input = torch.randn(1, self.NUM_FEATURES + self.NUM_ACTIONS) self.relaxed_sarsa_model = BoundedModule(self.sarsa_model, dummy_input) def setup_stores(self, store): # Logging setup self.store = store self.store.add_table('optimization', { 'mean_reward':float, 'final_value_loss':float, 'mean_std':float }) if self.advanced_logging: paper_constraint_cols = { 'avg_kl':float, 'max_kl':float, 'max_ratio':float, 'opt_step':int } value_cols = { 'heldout_gae_loss':float, 'heldout_returns_loss':float, 'train_gae_loss':float, 'train_returns_loss':float } weight_cols = {} for name, _ in self.policy_model.named_parameters(): name += "." for k in ["l1", "l2", "linf", "delta_l1", "delta_l2", "delta_linf"]: weight_cols[name + k] = float self.store.add_table('paper_constraints_train', paper_constraint_cols) self.store.add_table('paper_constraints_heldout', paper_constraint_cols) self.store.add_table('value_data', value_cols) self.store.add_table('weight_updates', weight_cols) if self.params.MODE == 'robust_ppo': robust_cols ={ 'eps': float, 'beta': float, 'kl': float, 'surrogate': float, 'entropy': float, 'loss': float, } self.store.add_table('robust_ppo_data', robust_cols) def __getattr__(self, x): ''' Allows accessing self.A instead of self.params.A ''' if x == 'params': return {} try: return getattr(self.params, x) except KeyError: raise AttributeError(x) def advantage_and_return(self, rewards, values, not_dones): """ Calculate GAE advantage, discounted returns, and true reward (average reward per trajectory) GAE: delta_t^V = r_t + discount * V(s_{t+1}) - V(s_t) using formula from John Schulman's code: V(s_t+1) = {0 if s_t is terminal {v_s_{t+1} if s_t not terminal and t != T (last step) {v_s if s_t not terminal and t == T """ assert shape_equal_cmp(rewards, values, not_dones) V_s_tp1 = ch.cat([values[:,1:], values[:, -1:]], 1) * not_dones deltas = rewards + self.GAMMA * V_s_tp1 - values # now we need to discount each path by gamma * lam advantages = ch.zeros_like(rewards) returns = ch.zeros_like(rewards) indices = get_path_indices(not_dones) for agent, start, end in indices: advantages[agent, start:end] = discount_path( \ deltas[agent, start:end], self.LAMBDA*self.GAMMA) returns[agent, start:end] = discount_path( \ rewards[agent, start:end], self.GAMMA) return advantages.clone().detach(), returns.clone().detach() def reset_envs(self, envs): ''' Resets environments and returns initial state with shape: (# actors, 1, ... state_shape) ''' if self.CPU: return cpu_tensorize([env.reset() for env in envs]).unsqueeze(1) else: return cu_tensorize([env.reset() for env in envs]).unsqueeze(1) def multi_actor_step(self, actions, envs): ''' Simulate a "step" by several actors on their respective environments Inputs: - actions, list of actions to take - envs, list of the environments in which to take the actions Returns: - completed_episode_info, a variable-length list of final rewards and episode lengths for the actors which have completed - rewards, a actors-length tensor with the rewards collected - states, a (actors, ... state_shape) tensor with resulting states - not_dones, an actors-length tensor with 0 if terminal, 1 otw ''' normed_rewards, states, not_dones = [], [], [] completed_episode_info = [] for action, env in zip(actions, envs): gym_action = action[0].cpu().numpy() new_state, normed_reward, is_done, info = env.step(gym_action) if is_done: completed_episode_info.append(info['done']) new_state = env.reset() # Aggregate normed_rewards.append([normed_reward]) not_dones.append([int(not is_done)]) states.append([new_state]) tensor_maker = cpu_tensorize if self.CPU else cu_tensorize data = list(map(tensor_maker, [normed_rewards, states, not_dones])) return [completed_episode_info, *data] def run_trajectories(self, num_saps, return_rewards=False, should_tqdm=False): """ Resets environments, and runs self.T steps in each environment in self.envs. If an environment hits a terminal state, the env is restarted and the terminal timestep marked. Each item in the tuple is a tensor in which the first coordinate represents the actor, and the second coordinate represents the time step. The third+ coordinates, if they exist, represent additional information for each time step. Inputs: None Returns: - rewards: (# actors, self.T) - not_dones: (# actors, self.T) 1 in timestep if terminal state else 0 - actions: (# actors, self.T, ) indices of actions - action_logprobs: (# actors, self.T, ) log probabilities of each action - states: (# actors, self.T, ... state_shape) states """ # Arrays to be updated with historic info envs = self.envs initial_states = self.reset_envs(envs) # Holds information (length and true reward) about completed episodes completed_episode_info = [] traj_length = int(num_saps // self.NUM_ACTORS) shape = (self.NUM_ACTORS, traj_length) all_zeros = [ch.zeros(shape) for i in range(3)] rewards, not_dones, action_log_probs = all_zeros actions_shape = shape + (self.NUM_ACTIONS,) actions = ch.zeros(actions_shape) # Mean of the action distribution. Used for avoid unnecessary recomputation. action_means = ch.zeros(actions_shape) # Log Std of the action distribution. action_stds = ch.zeros(actions_shape) states_shape = (self.NUM_ACTORS, traj_length+1) + initial_states.shape[2:] states = ch.zeros(states_shape) iterator = range(traj_length) if not should_tqdm else tqdm.trange(traj_length) assert self.NUM_ACTORS == 1 states[:, 0, :] = initial_states last_states = states[:, 0, :] for t in iterator: # assert shape_equal([self.NUM_ACTORS, self.NUM_FEATURES], last_states) # Retrieve probabilities # action_pds: (# actors, # actions), prob dists over actions # next_actions: (# actors, 1), indices of actions # next_action_probs: (# actors, 1), prob of taken actions last_states = self.apply_attack(last_states) # Note that for adversarial training, we use the state under perturbation to get the actions. # However in the trajectory we still save the state without perturbation as the true environment states are not perturbed. action_pds = self.policy_model(last_states) next_action_means, next_action_stds = action_pds next_actions = self.policy_model.sample(action_pds) next_action_log_probs = self.policy_model.get_loglikelihood(action_pds, next_actions) next_action_log_probs = next_action_log_probs.unsqueeze(1) # shape_equal([self.NUM_ACTORS, 1], next_action_log_probs) # if discrete, next_actions is (# actors, 1) # otw if continuous (# actors, 1, action dim) next_actions = next_actions.unsqueeze(1) # if self.policy_model.discrete: # assert shape_equal([self.NUM_ACTORS, 1], next_actions) # else: # assert shape_equal([self.NUM_ACTORS, 1, self.policy_model.action_dim]) ret = self.multi_actor_step(next_actions, envs) # done_info = List of (length, reward) pairs for each completed trajectory # (next_rewards, next_states, next_dones) act like multi-actor env.step() done_info, next_rewards, next_states, next_not_dones = ret # assert shape_equal([self.NUM_ACTORS, 1], next_rewards, next_not_dones) # assert shape_equal([self.NUM_ACTORS, 1, self.NUM_FEATURES], next_states) # If some of the actors finished AND this is not the last step # OR some of the actors finished AND we have no episode information if len(done_info) > 0 and (t != self.T - 1 or len(completed_episode_info) == 0): completed_episode_info.extend(done_info) # Update histories # each shape: (nact, t, ...) -> (nact, t + 1, ...) pairs = [ (rewards, next_rewards), (not_dones, next_not_dones), (actions, next_actions), # The sampled actions. (action_means, next_action_means), # The Gaussian mean of actions. # (action_stds, next_action_stds), # The Gaussian std of actions, is a constant, no need to save. (action_log_probs, next_action_log_probs), (states, next_states), ] last_states = next_states[:, 0, :] for total, v in pairs: if total is states: # Next states, stores in the next position. total[:, t+1] = v else: # The current action taken, and reward received. total[:, t] = v # Calculate the average episode length and true rewards over all the trajectories infos = np.array(list(zip(*completed_episode_info))) # print(infos) if infos.size > 0: _, ep_rewards = infos avg_episode_length, avg_episode_reward = np.mean(infos, axis=1) else: ep_rewards = [-1] avg_episode_length = -1 avg_episode_reward = -1 # Last state is never acted on, discard states = states[:,:-1,:] trajs = Trajectories(rewards=rewards, action_log_probs=action_log_probs, not_dones=not_dones, actions=actions, states=states, action_means=action_means, action_std=next_action_stds) to_ret = (avg_episode_length, avg_episode_reward, trajs) if return_rewards: to_ret += (ep_rewards,) return to_ret """Conduct adversarial attack using value network.""" def apply_attack(self, last_states): if self.params.ATTACK_RATIO < random.random(): # Only attack a portion of steps. return last_states eps = self.params.ATTACK_EPS if eps == "same": eps = self.params.ROBUST_PPO_EPS else: eps = float(eps) steps = self.params.ATTACK_STEPS if self.params.ATTACK_METHOD == "critic": # Find a state that is close the last_states and decreases value most. if steps > 0: if self.params.ATTACK_STEP_EPS == "auto": step_eps = eps / steps else: step_eps = float(self.params.ATTACK_STEP_EPS) clamp_min = last_states - eps clamp_max = last_states + eps # Random start. noise = torch.empty_like(last_states).uniform_(-step_eps, step_eps) states = last_states + noise with torch.enable_grad(): for i in range(steps): states = states.clone().detach().requires_grad_() value = self.val_model(states).mean(dim=1) value.backward() update = states.grad.sign() * step_eps # Clamp to +/- eps. states.data = torch.min(torch.max(states.data - update, clamp_min), clamp_max) self.val_model.zero_grad() return states.detach() else: return last_states elif self.params.ATTACK_METHOD == "random": # Apply an uniform random noise. noise = torch.empty_like(last_states).uniform_(-eps, eps) return (last_states + noise).detach() elif self.params.ATTACK_METHOD == "action": if steps > 0: if self.params.ATTACK_STEP_EPS == "auto": step_eps = eps / steps else: step_eps = float(self.params.ATTACK_STEP_EPS) clamp_min = last_states - eps clamp_max = last_states + eps # SGLD noise factor. We simply set beta=1. noise_factor = np.sqrt(2 * step_eps) noise = torch.randn_like(last_states) * noise_factor # The first step has gradient zero, so add the noise and projection directly. states = last_states + noise.sign() * step_eps # Current action at this state. old_action, old_stdev = self.policy_model(last_states) # Normalize stdev, avoid numerical issue old_stdev /= (old_stdev.mean()) old_action = old_action.detach() with torch.enable_grad(): for i in range(steps): states = states.clone().detach().requires_grad_() action_change = (self.policy_model(states)[0] - old_action) / old_stdev action_change = (action_change * action_change).sum(dim=1) action_change.backward() # Reduce noise at every step. noise_factor = np.sqrt(2 * step_eps) / (i+2) # Project noisy gradient to step boundary. update = (states.grad + noise_factor * torch.randn_like(last_states)).sign() * step_eps # Clamp to +/- eps. states.data = torch.min(torch.max(states.data + update, clamp_min), clamp_max) self.policy_model.zero_grad() return states.detach() else: return last_states elif self.params.ATTACK_METHOD == "sarsa" or self.params.ATTACK_METHOD == "sarsa+action": # Attack using a learned value network. assert self.params.ATTACK_SARSA_NETWORK is not None use_action = self.params.ATTACK_SARSA_ACTION_RATIO > 0 and self.params.ATTACK_METHOD == "sarsa+action" action_ratio = self.params.ATTACK_SARSA_ACTION_RATIO assert action_ratio >= 0 and action_ratio <= 1 if not hasattr(self, "sarsa_network"): self.sarsa_network = ValueDenseNet(state_dim=self.NUM_FEATURES+self.NUM_ACTIONS, init="normal") print("Loading sarsa network", self.params.ATTACK_SARSA_NETWORK) sarsa_ckpt = torch.load(self.params.ATTACK_SARSA_NETWORK) sarsa_meta = sarsa_ckpt['metadata'] sarsa_eps = sarsa_meta['sarsa_eps'] if 'sarsa_eps' in sarsa_meta else "unknown" sarsa_reg = sarsa_meta['sarsa_reg'] if 'sarsa_reg' in sarsa_meta else "unknown" sarsa_steps = sarsa_meta['sarsa_steps'] if 'sarsa_steps' in sarsa_meta else "unknown" print(f"Sarsa network was trained with eps={sarsa_eps}, reg={sarsa_reg}, steps={sarsa_steps}") if use_action: print(f"objective: {1.0 - action_ratio} * sarsa + {action_ratio} * action_change") else: print("Not adding action change objective.") self.sarsa_network.load_state_dict(sarsa_ckpt['state_dict']) if steps > 0: if self.params.ATTACK_STEP_EPS == "auto": step_eps = eps / steps else: step_eps = float(self.params.ATTACK_STEP_EPS) clamp_min = last_states - eps clamp_max = last_states + eps # Random start. noise = torch.empty_like(last_states).uniform_(-step_eps, step_eps) states = last_states + noise if use_action: # Current action at this state. old_action, old_stdev = self.policy_model(last_states) old_stdev /= (old_stdev.mean()) old_action = old_action.detach() with torch.enable_grad(): for i in range(steps): states = states.clone().detach().requires_grad_() # This is the mean action... actions = self.policy_model(states)[0] value = self.sarsa_network(torch.cat((last_states, actions), dim=1)).mean(dim=1) if use_action: action_change = (actions - old_action) / old_stdev # We want to maximize the action change, thus the minus sign. action_change = -(action_change * action_change).mean(dim=1) loss = action_ratio * action_change + (1.0 - action_ratio) * value else: action_change = 0.0 loss = value loss.backward() update = states.grad.sign() * step_eps # Clamp to +/- eps. states.data = torch.min(torch.max(states.data - update, clamp_min), clamp_max) self.val_model.zero_grad() return states.detach() else: return last_states elif self.params.ATTACK_METHOD == "none": return last_states else: raise ValueError(f'Unknown attack method {self.params.ATTACK_METHOD}') """Run trajectories and return saps and values for each state.""" def collect_saps(self, num_saps, should_log=True, return_rewards=False, should_tqdm=False, test=False): with torch.no_grad(): # Run trajectories, get values, estimate advantage output = self.run_trajectories(num_saps, return_rewards=return_rewards, should_tqdm=should_tqdm) if not return_rewards: avg_ep_length, avg_ep_reward, trajs = output else: avg_ep_length, avg_ep_reward, trajs, ep_rewards = output # No need to compute advantage function for testing. if not test: # If we are sharing weights between the policy network and # value network, we use the get_value function of the # *policy* to # estimate the value, instead of using the value # net if not self.SHARE_WEIGHTS: values = self.val_model(trajs.states).squeeze(-1) else: values = self.policy_model.get_value(trajs.states).squeeze(-1) # Calculate advantages and returns advantages, returns = self.advantage_and_return(trajs.rewards, values, trajs.not_dones) trajs.advantages = advantages trajs.returns = returns trajs.values = values assert shape_equal_cmp(trajs.advantages, trajs.returns, trajs.values) # Logging if should_log: msg = "Current mean reward: %f | mean episode length: %f" print(msg % (avg_ep_reward, avg_ep_length)) if not test: self.store.log_table_and_tb('optimization', { 'mean_reward': avg_ep_reward }) # Unroll the trajectories (actors, T, ...) -> (actors*T, ...) saps = trajs.unroll() to_ret = (saps, avg_ep_reward, avg_ep_length) if return_rewards: to_ret += (ep_rewards,) return to_ret def sarsa_steps(self, saps): # Begin advanged logging code assert saps.unrolled loss = torch.nn.SmoothL1Loss() action_std = torch.exp(self.policy_model.log_stdev).detach().requires_grad_(False) # Avoid backprop twice. # We treat all value epochs as one epoch. self.sarsa_eps_scheduler.set_epoch_length(self.params.VAL_EPOCHS * self.params.NUM_MINIBATCHES) self.sarsa_beta_scheduler.set_epoch_length(self.params.VAL_EPOCHS * self.params.NUM_MINIBATCHES) # We count from 1. self.sarsa_eps_scheduler.step_epoch() self.sarsa_beta_scheduler.step_epoch() # saps contains state->action->reward and not_done. for i in range(self.params.VAL_EPOCHS): # Create minibatches with shuffuling state_indices = np.arange(saps.rewards.nelement()) np.random.shuffle(state_indices) splits = np.array_split(state_indices, self.params.NUM_MINIBATCHES) # Minibatch SGD for selected in splits: def sel(*args): return [v[selected] for v in args] self.sarsa_opt.zero_grad() sel_states, sel_actions, sel_rewards, sel_not_dones = sel(saps.states, saps.actions, saps.rewards, saps.not_dones) self.sarsa_eps_scheduler.step_batch() self.sarsa_beta_scheduler.step_batch() inputs = torch.cat((sel_states, sel_actions), dim=1) # action_diff = self.sarsa_eps_scheduler.get_eps() * action_std # inputs_lb = torch.cat((sel_states, sel_actions - action_diff), dim=1).detach().requires_grad_(False) # inputs_ub = torch.cat((sel_states, sel_actions + action_diff), dim=1).detach().requires_grad_(False) # bounded_inputs = BoundedTensor(inputs, ptb=PerturbationLpNorm(norm=np.inf, eps=None, x_L=inputs_lb, x_U=inputs_ub)) bounded_inputs = BoundedTensor(inputs, ptb=PerturbationLpNorm(norm=np.inf, eps=self.sarsa_eps_scheduler.get_eps())) q = self.relaxed_sarsa_model(bounded_inputs).squeeze(-1) q_old = q[:-1] q_next = q[1:] * self.GAMMA * sel_not_dones[:-1] + sel_rewards[:-1] q_next = q_next.detach() # q_loss = (q_old - q_next).pow(2).sum(dim=-1).mean() q_loss = loss(q_old, q_next) # Compute the robustness regularization. if self.sarsa_eps_scheduler.get_eps() > 0 and self.params.SARSA_REG > 0: beta = self.sarsa_beta_scheduler.get_eps() ilb, iub = self.relaxed_sarsa_model.compute_bounds(IBP=True, method=None) if beta < 1: clb, cub = self.relaxed_sarsa_model.compute_bounds(IBP=False, method='backward') lb = beta * ilb + (1 - beta) * clb ub = beta * iub + (1 - beta) * cub else: lb = ilb ub = iub # Output dimension is 1. Remove the extra dimension and keep only the batch dimension. lb = lb.squeeze(-1) ub = ub.squeeze(-1) diff = torch.max(ub - q, q - lb) reg_loss = self.params.SARSA_REG * (diff * diff).mean() sarsa_loss = q_loss + reg_loss reg_loss = reg_loss.item() else: reg_loss = 0.0 sarsa_loss = q_loss sarsa_loss.backward() self.sarsa_opt.step() print(f'q_loss={q_loss.item():.6g}, reg_loss={reg_loss:.6g}, sarsa_loss={sarsa_loss.item():.6g}') if self.ANNEAL_LR: self.sarsa_scheduler.step() # print('value:', self.val_model(saps.states).mean().item()) return q_loss, q.mean() def take_steps(self, saps, logging=True, value_only=False): # Begin advanged logging code assert saps.unrolled should_adv_log = self.advanced_logging and \ self.n_steps % self.log_every == 0 and logging self.params.SHOULD_LOG_KL = self.advanced_logging and \ self.KL_APPROXIMATION_ITERS != -1 and \ self.n_steps % self.KL_APPROXIMATION_ITERS == 0 store_to_pass = self.store if should_adv_log else None # End logging code if should_adv_log: # collect some extra trajactory for validation of KL and max KL. num_saps = saps.advantages.shape[0] val_saps = self.collect_saps(num_saps, should_log=False)[0] out_train = self.policy_model(saps.states) out_val = self.policy_model(val_saps.states) old_pds = select_prob_dists(out_train, detach=True) val_old_pds = select_prob_dists(out_val, detach=True) # Update the value function before unrolling the trajectories # Pass the logging data into the function if applicable val_loss = ch.tensor(0.0) if not self.SHARE_WEIGHTS: val_loss = value_step(saps.states, saps.returns, saps.advantages, saps.not_dones, self.val_model, self.val_opt, self.params, store_to_pass).mean() if self.ANNEAL_LR: self.VALUE_SCHEDULER.step() if value_only: # Run the value iteration only. Return now. return val_loss if logging: self.store.log_table_and_tb('optimization', { 'final_value_loss': val_loss }) if self.MODE == 'robust_ppo' and logging: # Logging Robust PPO KL, entropy, etc. store_to_pass = self.store # Take optimizer steps args = [saps.states, saps.actions, saps.action_log_probs, saps.rewards, saps.returns, saps.not_dones, saps.advantages, self.policy_model, self.params, store_to_pass, self.n_steps] if self.MODE == 'robust_ppo' and isinstance(self.policy_model, CtsPolicy): args += [self.relaxed_policy_model, self.robust_eps_scheduler, self.robust_beta_scheduler] self.MAX_KL += self.MAX_KL_INCREMENT if should_adv_log: # Save old parameter to investigate weight updates. old_parameter = copy.deepcopy(self.policy_model.state_dict()) # Policy optimization step surr_loss = self.policy_step(*args).mean() # If the anneal_lr option is set, then we decrease the # learning rate at each training step if self.ANNEAL_LR: self.POLICY_SCHEDULER.step() if should_adv_log: log_value_losses(self, val_saps, 'heldout') log_value_losses(self, saps, 'train') old_pds = saps.action_means, saps.action_std paper_constraints_logging(self, saps, old_pds, table='paper_constraints_train') paper_constraints_logging(self, val_saps, val_old_pds, table='paper_constraints_heldout') log_weight_updates(self, old_parameter, self.policy_model.state_dict()) self.store['paper_constraints_train'].flush_row() self.store['paper_constraints_heldout'].flush_row() self.store['value_data'].flush_row() self.store['weight_updates'].flush_row() if self.params.MODE == 'robust_ppo': self.store['robust_ppo_data'].flush_row() return surr_loss, val_loss def train_step(self): ''' Take a training step, by first collecting rollouts, then calculating advantages, then taking a policy gradient step, and finally taking a value function step. Inputs: None Returns: - The current reward from the policy (per actor) ''' print("-" * 80) start_time = time.time() num_saps = self.T * self.NUM_ACTORS saps, avg_ep_reward, avg_ep_length = self.collect_saps(num_saps) surr_loss, val_loss = self.take_steps(saps) # Logging code print("Surrogate Loss:", surr_loss.item(), "| Value Loss:", val_loss.item()) print("Time elapsed (s):", time.time() - start_time) if not self.policy_model.discrete: mean_std = ch.exp(self.policy_model.log_stdev).mean() print("Agent stdevs: %s" % mean_std) self.store.log_table_and_tb('optimization', { 'mean_std': mean_std }) else: self.store['optimization'].update_row({ 'mean_std':np.nan }) self.store['optimization'].flush_row() sys.stdout.flush() sys.stderr.flush() # End logging code self.n_steps += 1 return avg_ep_reward def sarsa_step(self): ''' Take a training step, by first collecting rollouts, and taking a value function step. Inputs: None Returns: - The current reward from the policy (per actor) ''' print("-" * 80) start_time = time.time() num_saps = self.T * self.NUM_ACTORS saps, avg_ep_reward, avg_ep_length = self.collect_saps(num_saps, should_log=True, test=True) sarsa_loss, q = self.sarsa_steps(saps) print("Sarsa Loss:", sarsa_loss.item()) print("Q:", q.item()) print("Time elapsed (s):", time.time() - start_time) sys.stdout.flush() sys.stderr.flush() self.n_steps += 1 return avg_ep_reward def run_test(self, max_len=2048, compute_bounds=False, use_full_backward=False): print("-" * 80) start_time = time.time() if compute_bounds: self.create_relaxed_model() #saps, avg_ep_reward, avg_ep_length = self.collect_saps(num_saps=None, should_log=True, test=True, num_episodes=num_episodes) with torch.no_grad(): output = self.run_test_trajectories(max_len=max_len) ep_length, ep_reward, actions, action_means, states = output msg = "Episode reward: %f | episode length: %f" print(msg % (ep_reward, ep_length)) if compute_bounds: eps = float(self.params.ROBUST_PPO_EPS) if self.params.ATTACK_EPS == "same" else float(self.params.ATTACK_EPS) kl_upper_bound = get_state_kl_bound(self.relaxed_policy_model, states, action_means, eps=eps, beta=0.0, stdev=self.policy_model.log_stdev, use_full_backward=use_full_backward).mean() kl_upper_bound = kl_upper_bound.item() else: kl_upper_bound = float("nan") # Unroll the trajectories (actors, T, ...) -> (actors*T, ...) return ep_length, ep_reward, actions.cpu().numpy(), action_means.cpu().numpy(), states.cpu().numpy(), kl_upper_bound def run_test_trajectories(self, max_len, should_tqdm=False): # Arrays to be updated with historic info envs = self.envs initial_states = self.reset_envs(envs) # Holds information (length and true reward) about completed episodes completed_episode_info = [] shape = (1, max_len) rewards = ch.zeros(shape) actions_shape = shape + (self.NUM_ACTIONS,) actions = ch.zeros(actions_shape) # Mean of the action distribution. Used for avoid unnecessary recomputation. action_means = ch.zeros(actions_shape) states_shape = (1, max_len+1) + initial_states.shape[2:] states = ch.zeros(states_shape) iterator = range(max_len) if not should_tqdm else tqdm.trange(max_len) states[:, 0, :] = initial_states last_states = states[:, 0, :] for t in iterator: if (t+1) % 100 == 0: print('Step {} '.format(t+1)) # assert shape_equal([self.NUM_ACTORS, self.NUM_FEATURES], last_states) # Retrieve probabilities # action_pds: (# actors, # actions), prob dists over actions # next_actions: (# actors, 1), indices of actions maybe_attacked_last_states = self.apply_attack(last_states) action_pds = self.policy_model(maybe_attacked_last_states) next_action_means, next_action_stds = action_pds # Double check if the attack is within eps range. if self.params.ATTACK_METHOD != "none": max_eps = (maybe_attacked_last_states - last_states).abs().max() attack_eps = float(self.params.ROBUST_PPO_EPS) if self.params.ATTACK_EPS == "same" else float(self.params.ATTACK_EPS) if max_eps > attack_eps + 1e-5: raise RuntimeError(f"{max_eps} > {self.params.ATTACK_EPS}. Attack implementation has bug and eps is not correctly handled.") next_actions = self.policy_model.sample(action_pds) # if discrete, next_actions is (# actors, 1) # otw if continuous (# actors, 1, action dim) next_actions = next_actions.unsqueeze(1) ret = self.multi_actor_step(next_actions, envs) # done_info = List of (length, reward) pairs for each completed trajectory # (next_rewards, next_states, next_dones) act like multi-actor env.step() done_info, next_rewards, next_states, next_not_dones = ret # Update histories # each shape: (nact, t, ...) -> (nact, t + 1, ...) pairs = [ (rewards, next_rewards), (actions, next_actions), # The sampled actions. (action_means, next_action_means), # The sampled actions. (states, next_states), ] last_states = next_states[:, 0, :] for total, v in pairs: if total is states: # Next states, stores in the next position. total[:, t+1] = v else: # The current action taken, and reward received. total[:, t] = v # If some of the actors finished AND this is not the last step # OR some of the actors finished AND we have no episode information if len(done_info) > 0: completed_episode_info.extend(done_info) break if len(completed_episode_info) > 0: ep_length, ep_reward = completed_episode_info[0] else: ep_length = np.nan ep_reward = np.nan actions = actions[0][:t+1] action_means = action_means[0][:t+1] states = states[0][:t+1] to_ret = (ep_length, ep_reward, actions, action_means, states) return to_ret @staticmethod def agent_from_data(store, row, cpu, extra_params=None, override_params=None, excluded_params=None): ''' Initializes an agent from serialized data (via cox) Inputs: - store, the name of the store where everything is logged - row, the exact row containing the desired data for this agent - cpu, True/False whether to use the CPU (otherwise sends to GPU) - extra_params, a dictionary of extra agent parameters. Only used when a key does not exist from the loaded cox store. - override_params, a dictionary of agent parameters that will override current agent parameters. - excluded_params, a dictionary of parameters that we do not copy or override. Outputs: - agent, a constructed agent with the desired initialization and parameters - agent_params, the parameters that the agent was constructed with ''' ckpts = store['final_results'] get_item = lambda x: list(row[x])[0] items = ['val_model', 'policy_model', 'val_opt', 'policy_opt'] names = {i: get_item(i) for i in items} param_keys = list(store['metadata'].df.columns) param_values = list(store['metadata'].df.iloc[0,:]) def process_item(v): try: return v.item() except: return v param_values = [process_item(v) for v in param_values] agent_params = {k:v for k, v in zip(param_keys, param_values)} if 'adam_eps' not in agent_params: agent_params['adam_eps'] = 1e-5 if 'cpu' not in agent_params: agent_params['cpu'] = cpu # Update extra params if they do not exist in current parameters. if extra_params is not None: for k in extra_params.keys(): if k not in agent_params and k not in excluded_params: print(f'adding key {k}={extra_params[k]}') agent_params[k] = extra_params[k] if override_params is not None: for k in override_params.keys(): if k not in excluded_params and override_params[k] is not None and override_params[k] != agent_params[k]: print(f'overwriting key {k}: old={agent_params[k]}, new={override_params[k]}') agent_params[k] = override_params[k] agent = Trainer.agent_from_params(agent_params) def load_state_dict(model, ckpt_name): mapper = ch.device('cuda:0') if not cpu else ch.device('cpu') state_dict = ckpts.get_state_dict(ckpt_name, map_location=mapper) model.load_state_dict(state_dict) load_state_dict(agent.policy_model, names['policy_model']) load_state_dict(agent.val_model, names['val_model']) if agent.ANNEAL_LR: agent.POLICY_SCHEDULER.last_epoch = get_item('iteration') agent.VALUE_SCHEDULER.last_epoch = get_item('iteration') load_state_dict(agent.POLICY_ADAM, names['policy_opt']) load_state_dict(agent.val_opt, names['val_opt']) agent.envs = ckpts.get_pickle(get_item('envs')) return agent, agent_params @staticmethod def agent_from_params(params, store=None): ''' Construct a trainer object given a dictionary of hyperparameters. Trainer is in charge of sampling trajectories, updating policy network, updating value network, and logging. Inputs: - params, dictionary of required hyperparameters - store, a cox.Store object if logging is enabled Outputs: - A Trainer object for training a PPO/TRPO agent ''' agent_policy = policy_net_with_name(params['policy_net_type']) agent_value = value_net_with_name(params['value_net_type']) advanced_logging = params['advanced_logging'] and store is not None log_every = params['log_every'] if store is not None else 0 if params['cpu']: torch.set_num_threads(1) p = Trainer(agent_policy, agent_value, params, store, log_every=log_every, advanced_logging=advanced_logging) return p
class RobustDeterministicActorCriticNet(nn.Module, BaseNet): def __init__(self, state_dim, action_dim, actor_network, critic_network, mini_batch_size, actor_opt_fn, critic_opt_fn, robust_params=None): super(RobustDeterministicActorCriticNet, self).__init__() if robust_params is None: robust_params = {} self.use_loss_fusion = robust_params.get('use_loss_fusion', False) # Use loss fusion to reduce complexity for convex relaxation. Default is False. self.use_full_backward = robust_params.get('use_full_backward', False) if self.use_loss_fusion: # Use auto_LiRPA to compute the L2 norm directly. self.fc_action = model_mlp_any_with_loss(state_dim, actor_network, action_dim) modules = self.fc_action._modules # Auto LiRPA wrapper self.fc_action = BoundedModule( self.fc_action, (torch.empty(size=(1, state_dim)), torch.empty(size=(1, action_dim))), device=Config.DEVICE) # self.fc_action._modules = modules for n in self.fc_action.nodes: # Find the tanh neuron in computational graph if isinstance(n, BoundTanh): self.fc_action_after_tanh = n self.fc_action_pre_tanh = n.inputs[0] break else: # Fully connected layer with [state_dim, 400, 300, action_dim] neurons and ReLU activation function self.fc_action = model_mlp_any(state_dim, actor_network, action_dim) # auto_lirpa wrapper self.fc_action = BoundedModule( self.fc_action, (torch.empty(size=(1, state_dim)), ), device=Config.DEVICE) # Fully connected layer with [state_dim + action_dim, 400, 300, 1] self.fc_critic = model_mlp_any(state_dim + action_dim, critic_network, 1) # auto_lirpa wrapper self.fc_critic = BoundedModule( self.fc_critic, (torch.empty(size=(1, state_dim + action_dim)), ), device=Config.DEVICE) self.actor_params = self.fc_action.parameters() self.critic_params = self.fc_critic.parameters() self.actor_opt = actor_opt_fn(self.actor_params) self.critic_opt = critic_opt_fn(self.critic_params) self.to(Config.DEVICE) # Create identity specification matrices self.actor_identity = torch.eye(action_dim).repeat(mini_batch_size,1,1).to(Config.DEVICE) self.critic_identity = torch.eye(1).repeat(mini_batch_size,1,1).to(Config.DEVICE) self.action_dim = action_dim self.state_dim = state_dim def forward(self, obs): phi = self.feature(obs) action = self.actor(phi) return action def feature(self, obs): # Not used, originally this is a feature extraction network return tensor(obs) def actor(self, phi): if self.use_loss_fusion: self.fc_action(phi, torch.zeros(size=phi.size()[:1] + (self.action_dim,), device=Config.DEVICE)) return self.fc_action_after_tanh.forward_value else: return torch.tanh(self.fc_action(phi, method_opt="forward")) # Obtain element-wise lower and upper bounds for actor network through convex relaxations. def actor_bound(self, phi_lb, phi_ub, beta=1.0, eps=None, norm=np.inf, upper=True, lower=True, phi = None, center = None): if self.use_loss_fusion: # Use loss fusion (not typically enabled) assert center is not None ptb = PerturbationLpNorm(norm=norm, eps=eps, x_L=phi_lb, x_U=phi_ub) x = BoundedTensor(phi, ptb) val = self.fc_action(x, center.detach()) ilb, iub = self.fc_action.compute_bounds(IBP=True, method=None) if beta > 1e-10: clb, cub = self.fc_action.compute_bounds(IBP=False, method="backward", bound_lower=False, bound_upper=True) ub = cub * beta + iub * (1.0 - beta) return ub else: return iub else: assert center is None # Invoke auto_LiRPA for convex relaxation. ptb = PerturbationLpNorm(norm=norm, eps=eps, x_L=phi_lb, x_U=phi_ub) x = BoundedTensor(phi, ptb) if self.use_full_backward: clb, cub = self.fc_action.compute_bounds(x=(x,), IBP=False, method="backward") return cub, clb else: ilb, iub = self.fc_action.compute_bounds(x=(x,), IBP=True, method=None) if beta > 1e-10: clb, cub = self.fc_action.compute_bounds(IBP=False, method="backward") ub = cub * beta + iub * (1.0 - beta) lb = clb * beta + ilb * (1.0 - beta) return ub, lb else: return iub, ilb def critic(self, phi, a): return self.fc_critic(torch.cat([phi, a], dim=1), method_opt="forward") # Obtain element-wise lower and upper bounds for critic network through convex relaxations. def critic_bound(self, phi_lb, phi_ub, a_lb, a_ub, beta=1.0, eps=None, phi=None, action=None, norm=np.inf, upper=True, lower=True): x_L = torch.cat([phi_lb, a_lb], dim=1) x_U = torch.cat([phi_ub, a_ub], dim=1) ptb = PerturbationLpNorm(norm=norm, eps=eps, x_L=x_L, x_U=x_U) x = BoundedTensor(torch.cat([phi, action], dim=1), ptb) ilb, iub = self.fc_critic.compute_bounds(x=(x,), IBP=True, method=None) if beta > 1e-10: clb, cub = self.fc_critic.compute_bounds(IBP=False, method="backward") ub = cub * beta + iub * (1.0 - beta) lb = clb * beta + ilb * (1.0 - beta) return ub, lb else: return iub, ilb def load_state_dict(self, state_dict, strict=True): action_dict = OrderedDict() critic_dict = OrderedDict() for k in state_dict.keys(): if 'action' in k: pos = k.find('.') + 1 action_dict[k[pos:]] = state_dict[k] if 'critic' in k: pos = k.find('.') + 1 critic_dict[k[pos:]] = state_dict[k] # loading actor and critic networks separtely. this is requried for auto lirpa. self.fc_action.load_state_dict(action_dict) self.fc_critic.load_state_dict(critic_dict) def state_dict(self): # save actor and critic networks separtely. this is requried for auto lirpa. action_state_dict = self.fc_action.state_dict() critic_state_dict = self.fc_critic.state_dict() network_state_dict = OrderedDict() for k,v in action_state_dict.items(): network_state_dict["fc_action."+k] = v for k,v in critic_state_dict.items(): network_state_dict["fc_critic."+k] = v return network_state_dict