Unable to achieve the predicted effect

[userxzx2004]

When I allow this command，

(diff) abc@abc:~/SafeDiffuser$ python scripts/plan_maze2d.py --config config.maze2d --dataset maze2d-umaze-v1 --logbase logs
pybullet build time: Jan 29 2025 23:19:57
[ utils/colab ] Warning: not importing colab dependencies
[ utils/setup ] Reading config: config.maze2d:maze2d_umaze_v1
[ utils/setup ] Using overrides | config: config.maze2d | dataset: maze2d_umaze_v1
[ utils/setup ] Found extras: ['--logbase', 'logs']
[ utils/setup ] Overriding config | logbase : logs --> logs
[ utils/setup ] Lazy fstring | diffusion_loadpath : f:diffusion/H{horizon}_T{n_diffusion_steps} --> diffusion/H128_T64
[ utils/setup ] Setting exp_name to: plans/release_H128_T64_LimitsNormalizer_b1_condFalse
[ utils/setup ] Saved args to logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/args.json
[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/dataset_config.pkl

[utils/config ] Config: <class 'diffuser.datasets.sequence.GoalDataset'>
env: maze2d-umaze-v1
horizon: 128
max_path_length: 40000
normalizer: LimitsNormalizer
preprocess_fns: ['maze2d_set_terminals']
use_padding: False

[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/render_config.pkl

[utils/config ] Config: <class 'diffuser.utils.rendering.Maze2dRenderer'>
env: maze2d-umaze-v1

[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/model_config.pkl

[utils/config ] Config: <class 'diffuser.models.temporal.TemporalUnet'>
cond_dim: 4
dim_mults: (1, 4, 8)
horizon: 128
transition_dim: 6

[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/diffusion_config.pkl

[utils/config ] Config: <class 'diffuser.models.diffusion.GaussianDiffusion'>
action_dim: 2
action_weight: 1
clip_denoised: True
horizon: 128
loss_discount: 1
loss_type: l2
loss_weights: None
n_timesteps: 64
observation_dim: 4
predict_epsilon: False

[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/trainer_config.pkl

[utils/config ] Config: <class 'diffuser.utils.training.Trainer'>
bucket: None
ema_decay: 0.995
gradient_accumulate_every: 2
label_freq: 40000
n_reference: 50
n_samples: 10
results_folder: logs/maze2d-umaze-v1/diffusion/H128_T64
sample_freq: 0
save_freq: 1000
save_parallel: False
train_batch_size: 32
train_lr: 0.0002

load datafile: 100%|████████████████████████████████████████████████████████| 8/8 [00:00<00:00, 31.23it/s]
[ utils/preprocessing ] Segmented maze2d-umaze-v1 | 1565 paths | min length: 96 | max length: 3993
[ datasets/buffer ] Finalized replay buffer | 1566 episodes
[ datasets/buffer ] Fields:
actions: (1566, 40000, 2)
infos/goal: (1566, 40000, 2)
infos/qpos: (1566, 40000, 2)
infos/qvel: (1566, 40000, 2)
observations: (1566, 40000, 4)
rewards: (1566, 40000, 1)
terminals: (1566, 40000, 1)
timeouts: (1566, 40000, 1)
next_observations: (1566, 40000, 4)
normed_observations: (1566, 40000, 4)
normed_actions: (1566, 40000, 2)
[ models/temporal ] Channel dimensions: [(6, 32), (32, 128), (128, 256)]
[(6, 32), (32, 128), (128, 256)]

[ utils/serialization ] Loading model epoch: 1960000

step: 0 /100

/home/abc/anaconda3/envs/diff/lib/python3.8/site-packages/qpth/qp.py:83: UserWarning: torch.eig is deprecated in favor of torch.linalg.eig and will be removed in a future PyTorch release.
torch.linalg.eig returns complex tensors of dtype cfloat or cdouble rather than real tensors mimicking complex tensors.
L, _ = torch.eig(A)
should be replaced with
L_complex = torch.linalg.eigvals(A)
and
L, V = torch.eig(A, eigenvectors=True)
should be replaced with
L_complex, V_complex = torch.linalg.eig(A) (Triggered internally at ../aten/src/ATen/native/BatchLinearAlgebra.cpp:2910.)
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/0.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/0.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/1.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/2.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/3.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/4.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/5.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/6.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/7.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/8.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/9.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/10.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/11.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/12.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/13.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/14.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/15.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/16.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/17.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/18.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/19.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/20.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/21.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/22.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/23.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/24.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/25.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/26.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/27.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/28.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/29.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/30.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/31.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/32.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/33.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/34.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/35.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/36.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/37.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/38.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/39.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/40.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/41.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/42.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/43.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/44.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/45.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/46.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/47.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/48.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/49.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/50.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/51.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/52.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/53.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/54.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/55.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/56.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/57.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/58.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/59.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/60.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/61.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/62.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/63.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/64.png
elbo mean: 0.0
elbo std: 0.0
safe1: 0.4964413
safe2: 0.31419468
score mean: 0.9720310122454894
score std: 0.0
computation time: 2.7477188110351562
success rate: 1

/home/abc/SafeDiffuser/logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png

this is diffusion.py

[Weixy21]

In your case, you should also change the corresponding CBFs within SafeDiffuser according to your obstacle settings.

[userxzx2004]

in this select
x = self.invariance_umaze(x, xp1) #9

[userxzx2004]

Thank you for your suggestion, I'll give it a try.

[userxzx2004]

Do you mean I need to modify the CBF (Continuous Change Filter) to suit my environment? But I think I'm using a standard environment, so do I really need to modify the CBF specifically? Or is my understanding incorrect? Please give me some advice, thank you.

[Weixy21]

The CBFs are defined according to the locations (as well as geometries if you use some other types) of your obstacles. When you move your obstacles in the environment, this change has to be reflected in the CBFs (i.e., the b(x) function).

[userxzx2004]

I'm very interested in the security measures you mentioned, but this is the first time I've encountered this method. Could you please give me some specific suggestions, such as which part of the code needs to be modified?

[userxzx2004]

import numpy as np
import torch
from torch import nn
import pdb
from torch.autograd import Variable
from qpth.qp import QPFunction, QPSolvers
import einops

import diffuser.utils as utils
from .helpers import (
cosine_beta_schedule,
extract,
apply_conditioning,
Losses,
)

def normal_kl(mean1, logvar1, mean2, logvar2):
"""
Compute the KL divergence between two gaussians.

Shapes are automatically broadcasted, so batches can be compared to
scalars, among other use cases.
"""
tensor = None
for obj in (mean1, logvar1, mean2, logvar2):
    if isinstance(obj, torch.Tensor):
        tensor = obj
        break
assert tensor is not None, "at least one argument must be a Tensor"

# Force variances to be Tensors. Broadcasting helps convert scalars to
# Tensors, but it does not work for th.exp().
logvar1, logvar2 = [
    x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
    for x in (logvar1, logvar2)
]

return 0.5 * (
    -1.0
    + logvar2
    - logvar1
    + torch.exp(logvar1 - logvar2)
    + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
)

def approx_standard_normal_cdf(x):
"""
A fast approximation of the cumulative distribution function of the
standard normal.
"""
return 0.5 * (1.0 + torch.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * torch.pow(x, 3))))

def discretized_gaussian_log_likelihood(x, *, means, log_scales):
"""
Compute the log-likelihood of a Gaussian distribution discretizing to a
given image.

:param x: the target images. It is assumed that this was uint8 values,
          rescaled to the range [-1, 1].
:param means: the Gaussian mean Tensor.
:param log_scales: the Gaussian log stddev Tensor.
:return: a tensor like x of log probabilities (in nats).
"""
assert x.shape == means.shape == log_scales.shape
centered_x = x - means
inv_stdv = torch.exp(-log_scales)
plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
cdf_plus = approx_standard_normal_cdf(plus_in)
min_in = inv_stdv * (centered_x - 1.0 / 255.0)
cdf_min = approx_standard_normal_cdf(min_in)
log_cdf_plus = torch.log(cdf_plus.clamp(min=1e-12))
log_one_minus_cdf_min = torch.log((1.0 - cdf_min).clamp(min=1e-12))
cdf_delta = cdf_plus - cdf_min
log_probs = torch.where(
    x < -0.999,
    log_cdf_plus,
    torch.where(x > 0.999, log_one_minus_cdf_min, torch.log(cdf_delta.clamp(min=1e-12))),
)
assert log_probs.shape == x.shape
return log_probs

def mean_flat(tensor):
"""
Take the mean over all non-batch dimensions.
"""
return tensor.mean(dim=list(range(1, len(tensor.shape))))

class GaussianDiffusion(nn.Module):
def init(self, model, horizon, observation_dim, action_dim, n_timesteps=1000,
loss_type='l1', clip_denoised=False, predict_epsilon=True,
action_weight=1.0, loss_discount=1.0, loss_weights=None,
):
super().init()
self.horizon = horizon
self.observation_dim = observation_dim
self.action_dim = action_dim
self.transition_dim = observation_dim + action_dim
self.model = model
self.norm_mins = 0
self.norm_maxs = 0
self.safe1 = 0
self.safe2 = 0

    betas = cosine_beta_schedule(n_timesteps)
    alphas = 1. - betas
    alphas_cumprod = torch.cumprod(alphas, axis=0)
    alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]])

    self.n_timesteps = int(n_timesteps)
    self.clip_denoised = clip_denoised
    self.predict_epsilon = predict_epsilon

    self.register_buffer('betas', betas)
    self.register_buffer('alphas_cumprod', alphas_cumprod)
    self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)

    # calculations for diffusion q(x_t | x_{t-1}) and others
    self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
    self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))
    self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod))
    self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod))
    self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1))

    # calculations for posterior q(x_{t-1} | x_t, x_0)
    posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
    self.register_buffer('posterior_variance', posterior_variance)

    ## log calculation clipped because the posterior variance
    ## is 0 at the beginning of the diffusion chain
    self.register_buffer('posterior_log_variance_clipped',
        torch.log(torch.clamp(posterior_variance, min=1e-20)))
    self.register_buffer('posterior_mean_coef1',
        betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
    self.register_buffer('posterior_mean_coef2',
        (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))

    ## get loss coefficients and initialize objective
    loss_weights = self.get_loss_weights(action_weight, loss_discount, loss_weights)
    self.loss_fn = Losses[loss_type](loss_weights, self.action_dim)

def _format_conditions(self, conditions, batch_size):
    conditions = utils.apply_dict(
        self.normalizer.normalize,
        conditions,
        'observations',
    )
    conditions = utils.to_torch(conditions, dtype=torch.float32, device='cuda:0')
    conditions = utils.apply_dict(
        einops.repeat,
        conditions,
        'd -> repeat d', repeat=batch_size,
    )
    return conditions

def get_loss_weights(self, action_weight, discount, weights_dict):
    '''
        sets loss coefficients for trajectory

        action_weight   : float
            coefficient on first action loss
        discount   : float
            multiplies t^th timestep of trajectory loss by discount**t
        weights_dict    : dict
            { i: c } multiplies dimension i of observation loss by c
    '''
    self.action_weight = action_weight

    dim_weights = torch.ones(self.transition_dim, dtype=torch.float32)

    ## set loss coefficients for dimensions of observation
    if weights_dict is None: weights_dict = {}
    for ind, w in weights_dict.items():
        dim_weights[self.action_dim + ind] *= w

    ## decay loss with trajectory timestep: discount**t
    discounts = discount ** torch.arange(self.horizon, dtype=torch.float)
    discounts = discounts / discounts.mean()
    loss_weights = torch.einsum('h,t->ht', discounts, dim_weights)

    ## manually set a0 weight
    loss_weights[0, :self.action_dim] = action_weight
    return loss_weights

#------------------------------------------ sampling ------------------------------------------#

def predict_start_from_noise(self, x_t, t, noise):
    '''
        if self.predict_epsilon, model output is (scaled) noise;
        otherwise, model predicts x0 directly
    '''
    if self.predict_epsilon:
        return (
            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
            extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
        )
    else:
        return noise

def q_posterior(self, x_start, x_t, t):

    posterior_mean = (
        extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
        extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
    )
    posterior_variance = extract(self.posterior_variance, t, x_t.shape)
    posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
    return posterior_mean, posterior_variance, posterior_log_variance_clipped

def p_mean_variance(self, x, cond, t):

    x_recon = self.predict_start_from_noise(x, t=t, noise=self.model(x, cond, t))

    if self.clip_denoised:
        x_recon.clamp_(-1., 1.)
    else:
        assert RuntimeError()

    model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
            x_start=x_recon, x_t=x, t=t)
    return model_mean, posterior_variance, posterior_log_variance

@torch.no_grad()   #only for sampling
def Shield(self, x0, xp10):  #Truncate method

    x = x0.clone()
    xp1 = xp10.clone()

    xp1 = xp1.squeeze(0)

    nBatch = xp1.shape[0]

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    b = ((xp1[:,2:3] - off_y)/yr)**2 + ((xp1[:,3:4] - off_x)/xr)**2 - 1

    for k in range(nBatch):
        if b[k, 0] < 0: 
            theta = torch.atan2((xp1[k,2:3] - off_y)/yr, (xp1[k,3:4] - off_x)/xr)
            xp1[k,2] = yr*torch.sin(theta) + off_y
            xp1[k,3] = xr*torch.cos(theta) + off_x

    b = ((xp1[:,2:3] - off_y)/yr)**2 + ((xp1[:,3:4] - off_x)/xr)**2 - 1

     #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b2 = ((xp1[:,2:3] - off_y)/yr)**4 + ((xp1[:,3:4] - off_x)/xr)**4 - 1

    self.safe1 = torch.min(b[:,0])
    self.safe2 = torch.min(b2[:,0])

    xp1 = xp1.unsqueeze(0)
    return xp1

@torch.no_grad()   #only for sampling
def GD(self, x0, xp10):    #classifier guidance or potential-based method

    x = x0.clone()
    xp1 = xp10.clone()

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    b = ((xp1[:,2:3] - off_y)/yr)**2 + ((xp1[:,3:4] - off_x)/xr)**2 - 1

    #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b2 = ((xp1[:,2:3] - off_y)/yr)**4 + ((xp1[:,3:4] - off_x)/xr)**4 - 1

    for k in range(nBatch):
        if b[k, 0] < 0.1:  # 0, 0.2
            u1 = 0.2/(2*((xp1[k,2:3] - off_y)/yr)/yr)
            u2 = 0.2/(2*((xp1[k,3:4] - off_x)/xr)/xr)
            xp1[k,2] = xp1[k,2] + u1*0.001  #note no 0.1/0.01 for GD, but has for potential
            xp1[k,3] = xp1[k,3] + u2*0.001
        elif b2[k, 0] < 0.1:  # 0, 0.2
            u1 = 0.2/(4*((xp1[k,2:3] - off_y)/yr)**3/yr)
            u2 = 0.2/(4*((xp1[k,3:4] - off_x)/xr)**3/xr)
            xp1[k,2] = xp1[k,2] + u1*0.001
            xp1[k,3] = xp1[k,3] + u2*0.001
        # else:
        #     x[k,2] = xp1[k,2]
        #     x[k,3] = xp1[k,3]

    self.safe1 = torch.min(b[:,0])
    self.safe2 = torch.min(b2[:,0])

    xp1 = xp1.unsqueeze(0)
    return xp1

@torch.no_grad()   #only for sampling
def invariance_umaze(self, x, xp1):   #  RoS-diffuser for umaze

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1.52/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1.52/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(2.5-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2.5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

   #CBF
    b = 1 - ((x[:,2:3] - off_y)/yr)**4 - ((x[:,3:4] - off_x)/xr)**4
    Lfb = 0
    Lgbu1 = -4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = -4*((x[:,3:4] - off_x)/xr)**3/xr

    G = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    G = G.unsqueeze(1)
    k = 1
    h = Lfb + k*b

    self.safe1 = torch.min(b[:,0])

    #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1.2/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*0.6/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(2-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2.5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1
    Lfb = 0
    Lgbu1 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe2 = torch.min(b[:,0])

    G1 = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    G1 = G1.unsqueeze(1)
    k = 1
    h1 = Lfb + k*b

    G = torch.cat([G, G1], dim = 1)
    h = torch.cat([h, h1], dim = 1)
    

    q = -ref[:,2:4].to(G.device)
    Q = Variable(torch.eye(2))
    Q = Q.unsqueeze(0).expand(nBatch, 2, 2).to(G.device)
    
    e = Variable(torch.Tensor())
    out = QPFunction(verbose=-1, solver = QPSolvers.PDIPM_BATCHED)(Q, q, G, h, e, e)

    rt = xp1.clone()      
    rt[:,2:4] = x[:,2:4] + out
    rt = rt.unsqueeze(0)
    return rt

@torch.no_grad()   #only for sampling
def invariance_umaze_relax(self, x, xp1, t):  #  ReS-diffuser for umaze

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1.52/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1.52/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(2.5-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2.5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = 1 - ((x[:,2:3] - off_y)/yr)**4 - ((x[:,3:4] - off_x)/xr)**4
    Lfb = 0
    Lgbu1 = -4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = -4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe1 = torch.min(b[:,0])

    if t >= 10:
        sign = 100   #relax
    else:
        sign = 0   #non-relax

    rx0 = torch.zeros_like(Lgbu1).to(b.device)
    rx1 = sign*torch.ones_like(Lgbu1).to(b.device)

    G = torch.cat([-Lgbu1, -Lgbu2, rx1, rx0], dim = 1)
    G = G.unsqueeze(1)
    k = 1
    h = Lfb + k*b

    #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1.2/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*0.6/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(2-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2.5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1
    Lfb = 0
    Lgbu1 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe2 = torch.min(b[:,0])

    G1 = torch.cat([-Lgbu1, -Lgbu2, rx0, rx1], dim = 1)
    G1 = G1.unsqueeze(1)
    k = 1
    h1 = Lfb + k*b

    G = torch.cat([G, G1], dim = 1)
    h = torch.cat([h, h1], dim = 1)
    

    q = -ref[:,2:4].to(G.device)
    q0 = torch.zeros_like(q).to(G.device)
    q = torch.cat([q, q0], dim = 1)
    Q = Variable(torch.eye(4))
    Q = Q.unsqueeze(0).expand(nBatch, 4, 4).to(G.device)
    
    e = Variable(torch.Tensor())
    out = QPFunction(verbose=-1, solver = QPSolvers.PDIPM_BATCHED)(Q, q, G, h, e, e)

    rt = xp1.clone()      
    rt[:,2:4] = x[:,2:4] + out[:,0:2]
    rt = rt.unsqueeze(0)
    return rt

@torch.no_grad()   #only for sampling
def invariance(self, x, xp1):    #  RoS-diffuser for maze2d-large-v1

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**2 + ((x[:,3:4] - off_x)/xr)**2 - 1 - 0.01  # robust term 09/25
    Lfb = 0
    Lgbu1 = 2*((x[:,2:3] - off_y)/yr)/yr
    Lgbu2 = 2*((x[:,3:4] - off_x)/xr)/xr

    G = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    G = G.unsqueeze(1)
    k = 1
    h = Lfb + k*b

    self.safe1 = torch.min(b[:,0] + 0.01)  # robust term 09/25

    #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 - 0.01 # robust term 09/25
    Lfb = 0
    Lgbu1 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe2 = torch.min(b[:,0]+ 0.01) # robust term 09/25

    G1 = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    G1 = G1.unsqueeze(1)
    k = 1
    h1 = Lfb + k*b

    G = torch.cat([G, G1], dim = 1)
    h = torch.cat([h, h1], dim = 1)
    

    q = -ref[:,2:4].to(G.device)
    Q = Variable(torch.eye(2))
    Q = Q.unsqueeze(0).expand(nBatch, 2, 2).to(G.device)
    
    e = Variable(torch.Tensor())
    out = QPFunction(verbose=-1, solver = QPSolvers.PDIPM_BATCHED)(Q, q, G, h, e, e)

    rt = xp1.clone()      
    rt[:,2:4] = x[:,2:4] + out
    # print(out)
    rt = rt.unsqueeze(0)
    return rt

@torch.no_grad()   #only for sampling
def invariance_cf(self, x, xp1):  # closed form solution,  RoS-diffuser for maze2d-large-v1

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b0 = ((x[:,2:3] - off_y)/yr)**2 + ((x[:,3:4] - off_x)/xr)**2 - 1 - 0.01  # robust term 09/25
    Lfb = 0
    Lgbu1 = 2*((x[:,2:3] - off_y)/yr)/yr
    Lgbu2 = 2*((x[:,3:4] - off_x)/xr)/xr

    G0 = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    k = 1
    h0 = Lfb + k*b0

    self.safe1 = torch.min(b0[:,0] + 0.01)  # robust term 09/25

    #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 - 0.01 # robust term 09/25
    Lfb = 0
    Lgbu1 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe2 = torch.min(b[:,0]+ 0.01) # robust term 09/25

    G1 = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    k = 1
    h1 = Lfb + k*b
    
    q = -ref[:,2:4].to(b.device)
    
    y1_bar = 1*G0  # H or Q = identity matrix
    y2_bar = 1*G1
    u_bar = -1*q
    p1_bar = h0 - torch.sum(G0*u_bar,dim = 1).unsqueeze(1)
    p2_bar = h1 - torch.sum(G1*u_bar,dim = 1).unsqueeze(1)

    G = torch.cat([torch.sum(y1_bar*y1_bar,dim = 1).unsqueeze(1).unsqueeze(0), torch.sum(y1_bar*y2_bar,dim = 1).unsqueeze(1).unsqueeze(0), torch.sum(y2_bar*y1_bar,dim = 1).unsqueeze(1).unsqueeze(0), torch.sum(y2_bar*y2_bar,dim = 1).unsqueeze(1).unsqueeze(0)], dim = 0)
    #G = 1*[y1_bar*y1_bar', y1_bar*y2_bar'; y2_bar*y1_bar', y2_bar*y2_bar']
    w_p1_bar = torch.clamp(p1_bar, max=0)
    w_p2_bar = torch.clamp(p2_bar, max=0)

    # G 0-(1,1), 1-(1,2), 2-(2,1), 3-(2,2)
    lambda1 = torch.where(G[2]*w_p2_bar < G[3]*p1_bar, torch.zeros_like(p1_bar), torch.where(G[1]*w_p1_bar < G[0]*p2_bar, w_p1_bar/G[0], torch.clamp(G[3]*p1_bar - G[2]*p2_bar, max=0)/(G[0]*G[3] - G[1]*G[2])))
    
    lambda2 = torch.where(G[2]*w_p2_bar < G[3]*p1_bar, w_p2_bar/G[3], torch.where(G[1]*w_p1_bar < G[0]*p2_bar, torch.zeros_like(p1_bar), torch.clamp(G[0]*p2_bar - G[1]*p1_bar, max=0)/(G[0]*G[3] - G[1]*G[2])))

    out = lambda1*y1_bar + lambda2*y2_bar + u_bar
    rt = xp1.clone()      
    rt[:,2:4] = x[:,2:4] + out
    # print(out)
    rt = rt.unsqueeze(0)
    return rt
    



@torch.no_grad()   #only for sampling
def invariance_relax(self, x, xp1, t):  #  ReS-diffuser for maze2d-large-v1

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**2 + ((x[:,3:4] - off_x)/xr)**2 - 1 - 0.01
    Lfb = 0
    Lgbu1 = 2*((x[:,2:3] - off_y)/yr)/yr
    Lgbu2 = 2*((x[:,3:4] - off_x)/xr)/xr

    self.safe1 = torch.min(b[:,0] + 0.01)

    if t >= 10:   # debug  10
        sign = 100   #relax
    else:
        sign = 0   #non-relax

    rx0 = torch.zeros_like(Lgbu1).to(b.device)
    rx1 = sign*torch.ones_like(Lgbu1).to(b.device)

    G = torch.cat([-Lgbu1, -Lgbu2, rx1, rx0], dim = 1)
    G = G.unsqueeze(1)
    k = 1
    h = Lfb + k*b

    #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 - 0.01
    Lfb = 0
    Lgbu1 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe2 = torch.min(b[:,0] + 0.01)

    G1 = torch.cat([-Lgbu1, -Lgbu2, rx0, rx1], dim = 1)
    G1 = G1.unsqueeze(1)
    k = 1
    h1 = Lfb + k*b

    G = torch.cat([G, G1], dim = 1)
    h = torch.cat([h, h1], dim = 1)
    

    q = -ref[:,2:4].to(G.device)
    q0 = torch.zeros_like(q).to(G.device)
    q = torch.cat([q, q0], dim = 1)
    Q = Variable(torch.eye(4))
    Q = Q.unsqueeze(0).expand(nBatch, 4, 4).to(G.device)
    
    e = Variable(torch.Tensor())
    out = QPFunction(verbose=-1, solver = QPSolvers.PDIPM_BATCHED)(Q, q, G, h, e, e)

    rt = xp1.clone()      
    rt[:,2:4] = x[:,2:4] + out[:,0:2]
    rt = rt.unsqueeze(0)
    return rt

@torch.no_grad()   #only for sampling
def invariance_relax_cf(self, x, xp1, t):  # closed-form solution, ReS-diffuser for maze2d-large-v1

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**2 + ((x[:,3:4] - off_x)/xr)**2 - 1 - 0.01
    Lfb = 0
    Lgbu1 = 2*((x[:,2:3] - off_y)/yr)/yr
    Lgbu2 = 2*((x[:,3:4] - off_x)/xr)/xr

    self.safe1 = torch.min(b[:,0] + 0.01)

    if t >= 10:   # debug  10
        sign = 100   #relax
    else:
        sign = 0   #non-relax

    rx0 = torch.zeros_like(Lgbu1).to(b.device)
    rx1 = sign*torch.ones_like(Lgbu1).to(b.device)

    G0 = torch.cat([-Lgbu1, -Lgbu2, rx1, rx0], dim = 1)
    k = 1
    h0 = Lfb + k*b

    #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 - 0.01
    Lfb = 0
    Lgbu1 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe2 = torch.min(b[:,0] + 0.01)

    G1 = torch.cat([-Lgbu1, -Lgbu2, rx0, rx1], dim = 1)
    k = 1
    h1 = Lfb + k*b
    

    q = -ref[:,2:4].to(G0.device)
    q0 = torch.zeros_like(q).to(G0.device)
    q = torch.cat([q, q0], dim = 1)

    y1_bar = 1*G0  # H or Q = identity matrix
    y2_bar = 1*G1
    u_bar = -1*q
    p1_bar = h0 - torch.sum(G0*u_bar,dim = 1).unsqueeze(1)
    p2_bar = h1 - torch.sum(G1*u_bar,dim = 1).unsqueeze(1)

    G = torch.cat([torch.sum(y1_bar*y1_bar,dim = 1).unsqueeze(1).unsqueeze(0), torch.sum(y1_bar*y2_bar,dim = 1).unsqueeze(1).unsqueeze(0), torch.sum(y2_bar*y1_bar,dim = 1).unsqueeze(1).unsqueeze(0), torch.sum(y2_bar*y2_bar,dim = 1).unsqueeze(1).unsqueeze(0)], dim = 0)
    #G = 1*[y1_bar*y1_bar', y1_bar*y2_bar'; y2_bar*y1_bar', y2_bar*y2_bar']
    w_p1_bar = torch.clamp(p1_bar, max=0)
    w_p2_bar = torch.clamp(p2_bar, max=0)

    # G 0-(1,1), 1-(1,2), 2-(2,1), 3-(2,2)
    lambda1 = torch.where(G[2]*w_p2_bar < G[3]*p1_bar, torch.zeros_like(p1_bar), torch.where(G[1]*w_p1_bar < G[0]*p2_bar, w_p1_bar/G[0], torch.clamp(G[3]*p1_bar - G[2]*p2_bar, max=0)/(G[0]*G[3] - G[1]*G[2])))
    
    lambda2 = torch.where(G[2]*w_p2_bar < G[3]*p1_bar, w_p2_bar/G[3], torch.where(G[1]*w_p1_bar < G[0]*p2_bar, torch.zeros_like(p1_bar), torch.clamp(G[0]*p2_bar - G[1]*p1_bar, max=0)/(G[0]*G[3] - G[1]*G[2])))

    out = lambda1*y1_bar + lambda2*y2_bar + u_bar
    rt = xp1.clone()    
    rt[:,2:4] = x[:,2:4] + out[:,0:2]
    # print(out)
    rt = rt.unsqueeze(0)
    return rt


@torch.no_grad()   #only for sampling
def invariance_relax_narrow(self, x, xp1, t):  #  ReS-diffuser for maze2d-large-v1,  narrow passage case

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    if t >= 10:   # debug  10
        sign = 1   #relax
    else:
        sign = 0   #non-relax

    #normalize obstacle 1,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])

    off_x = 2*(5.5-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 - 0.4  # 0.01
    Lfb = 0
    Lgbu1 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    rx0 = torch.zeros_like(Lgbu1).to(b.device)
    rx1 = sign*torch.ones_like(Lgbu1).to(b.device)

    self.safe1 = torch.min(b[:,0] + 0.01)

    G1 = torch.cat([-Lgbu1, -Lgbu2, rx1, rx0, rx0, rx0, rx0, rx0], dim = 1)
    G1 = G1.unsqueeze(1)
    k = 1
    h1 = Lfb + k*b

    ########################################### obs 2
    off_x = 2*(5.5-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b2 = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 - 0.6 #0.01
    Lfb = 0
    Lgbu12 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu22 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe2 = torch.min(b2[:,0] + 0.01)

    G2 = torch.cat([-Lgbu12, -Lgbu22, rx0, rx1, rx0, rx0, rx0, rx0], dim = 1)
    G2 = G2.unsqueeze(1)
    k = 1
    h2 = Lfb + k*b2

    ########################################### obs 3
    off_x = 2*(3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2.5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b3 = ((x[:,2:3] - off_y)/yr/0.5)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 - 0.01
    Lfb = 0
    Lgbu13 = 4*((x[:,2:3] - off_y)/yr/0.5)**3/yr
    Lgbu23 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    G3 = torch.cat([-Lgbu13, -Lgbu23, rx0, rx0, rx1, rx0, rx0, rx0], dim = 1)
    G3 = G3.unsqueeze(1)
    k = 1
    h3 = Lfb + k*b3

    ########################################### obs 4
    off_x = 2*(8.5-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(3.5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b4 = ((x[:,2:3] - off_y)/yr/1.8)**4 + ((x[:,3:4] - off_x)/xr/1.8)**4 - 1 - 0.01
    Lfb = 0
    Lgbu14 = 4*((x[:,2:3] - off_y)/yr/1.8)**3/yr
    Lgbu24 = 4*((x[:,3:4] - off_x)/xr/1.8)**3/xr

    G4 = torch.cat([-Lgbu14, -Lgbu24, rx0, rx0, rx0, rx1, rx0, rx0], dim = 1)
    G4 = G4.unsqueeze(1)
    k = 1
    h4 = Lfb + k*b4

    ########################################### obs 5
    off_x = 2*(7.6-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(7-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b5 = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 - 0.4 #0.01
    Lfb = 0
    Lgbu15 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu25 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    G5 = torch.cat([-Lgbu15, -Lgbu25, rx0, rx0, rx0, rx0, rx1, rx0], dim = 1)
    G5 = G5.unsqueeze(1)
    k = 1
    h5 = Lfb + k*b5

    ########################################### obs 6
    off_x = 2*(10-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(6.3-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b6 = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 - 0.01
    Lfb = 0
    Lgbu16 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu26 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    G6 = torch.cat([-Lgbu16, -Lgbu26, rx0, rx0, rx0, rx0, rx0, rx1], dim = 1)
    G6 = G6.unsqueeze(1)
    k = 1
    h6 = Lfb + k*b6

    b0 = torch.cat([b, b2, b3, b4, b5, b6], dim = 1)
    idx = torch.argmin(b0, dim = 1).cpu().numpy()
    G0 = torch.cat([G1, G2, G3, G4, G5, G6], dim = 1)
    h0 = torch.cat([h1, h2, h3, h4, h5, h6], dim = 1)
    rows = len(G0[:,0,0])
    G = []
    h = []
    for i in range(rows):
        G.append(G0[i:i+1,idx[i]:idx[i]+1])
        h.append(h0[i:i+1,idx[i]:idx[i]+1])
    G = torch.cat(G, dim = 0)
    h = torch.cat(h, dim = 0)


    

    # G = torch.cat([G1, G2, G3, G4, G5, G6], dim = 1)
    # h = torch.cat([h1, h2, h3, h4, h5, h6], dim = 1)

    # G = torch.cat([G1, G2, G3, G5], dim = 1)
    # h = torch.cat([h1, h2, h3, h5], dim = 1)
    

    q = -ref[:,2:4].to(G.device)
    q0 = torch.zeros_like(q).to(G.device)
    q = torch.cat([q, q0, q0, q0], dim = 1)
    Q = Variable(torch.eye(8))
    Q = Q.unsqueeze(0).expand(nBatch, 8, 8).to(G.device)
    
    e = Variable(torch.Tensor())
    out = QPFunction(verbose=-1, solver = QPSolvers.PDIPM_BATCHED)(Q, q, G, h, e, e)

    rt = xp1.clone()      
    rt[:,2:4] = x[:,2:4] + out[:,0:2]
    rt = rt.unsqueeze(0)
    return rt


@torch.no_grad()   #only for sampling
def invariance_time(self, x, xp1, t):  #  TVS-diffuser for maze2d-large-v1
    t_bias = 5  #50

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**2 + ((x[:,3:4] - off_x)/xr)**2 - nn.Sigmoid()(t_bias - t) -0.01
    Lfb = nn.Sigmoid()(t_bias - t)*(1 - nn.Sigmoid()(t_bias - t))
    Lgbu1 = 2*((x[:,2:3] - off_y)/yr)/yr
    Lgbu2 = 2*((x[:,3:4] - off_x)/xr)/xr

    self.safe1 = torch.min(b[:,0] + 0.01)

    G = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    G = G.unsqueeze(1)
    k = 1  #0.3
    h = Lfb + k*b

    #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - nn.Sigmoid()(t_bias - t) - 0.01
    Lfb = nn.Sigmoid()(t_bias - t)*(1 - nn.Sigmoid()(t_bias - t))
    Lgbu1 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe2 = torch.min(b[:,0] + 0.01)

    G1 = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    G1 = G1.unsqueeze(1)
    k = 1  #0.4
    h1 = Lfb + k*b

    G = torch.cat([G, G1], dim = 1)
    h = torch.cat([h, h1], dim = 1)
    

    q = -ref[:,2:4].to(G.device)
    Q = Variable(torch.eye(2))
    Q = Q.unsqueeze(0).expand(nBatch, 2, 2).to(G.device)
    
    e = Variable(torch.Tensor())
    out = QPFunction(verbose=-1, solver = QPSolvers.PDIPM_BATCHED)(Q, q, G, h, e, e)

    rt = xp1.clone()      
    rt[:,2:4] = x[:,2:4] + out
    rt = rt.unsqueeze(0)
    return rt

@torch.no_grad()   #only for sampling
def invariance_time_cf(self, x, xp1, t):  # closed-form solution, TVS-diffuser for maze2d-large-v1
    t_bias = 5  #50 

    x = x.squeeze(0)
    xp1 = xp1.squeeze(0)

    nBatch = x.shape[0]
    ref = xp1 - x

    #normalize obstacle 1, x-1, y-0  x = 1/12*np.cos(theta) + 5.5/12, y = 1/9*np.sin(theta) + 5/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**2 + ((x[:,3:4] - off_x)/xr)**2 - nn.Sigmoid()(t_bias - t) -0.01
    Lfb = nn.Sigmoid()(t_bias - t)*(1 - nn.Sigmoid()(t_bias - t))
    Lgbu1 = 2*((x[:,2:3] - off_y)/yr)/yr
    Lgbu2 = 2*((x[:,3:4] - off_x)/xr)/xr

    self.safe1 = torch.min(b[:,0] + 0.01)

    G0 = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    k = 1  #0.3
    h0 = Lfb + k*b

    #normalize obstacle 2,  x = 1/12*np.sqrt(np.abs(np.cos(theta)))*np.sign(np.cos(theta)) + 5.3/12, y = 1/9*np.sqrt(np.abs(np.sin(theta)))*np.sign(np.sin(theta)) + 2/9
    xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1

    #CBF
    b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - nn.Sigmoid()(t_bias - t) - 0.01
    Lfb = nn.Sigmoid()(t_bias - t)*(1 - nn.Sigmoid()(t_bias - t))
    Lgbu1 = 4*((x[:,2:3] - off_y)/yr)**3/yr
    Lgbu2 = 4*((x[:,3:4] - off_x)/xr)**3/xr

    self.safe2 = torch.min(b[:,0] + 0.01)

    G1 = torch.cat([-Lgbu1, -Lgbu2], dim = 1)
    k = 1  #0.4
    h1 = Lfb + k*b
    

    q = -ref[:,2:4].to(G0.device)

    y1_bar = 1*G0  # H or Q = identity matrix
    y2_bar = 1*G1
    u_bar = -1*q
    p1_bar = h0 - torch.sum(G0*u_bar,dim = 1).unsqueeze(1)
    p2_bar = h1 - torch.sum(G1*u_bar,dim = 1).unsqueeze(1)

    G = torch.cat([torch.sum(y1_bar*y1_bar,dim = 1).unsqueeze(1).unsqueeze(0), torch.sum(y1_bar*y2_bar,dim = 1).unsqueeze(1).unsqueeze(0), torch.sum(y2_bar*y1_bar,dim = 1).unsqueeze(1).unsqueeze(0), torch.sum(y2_bar*y2_bar,dim = 1).unsqueeze(1).unsqueeze(0)], dim = 0)
    #G = 1*[y1_bar*y1_bar', y1_bar*y2_bar'; y2_bar*y1_bar', y2_bar*y2_bar']
    w_p1_bar = torch.clamp(p1_bar, max=0)
    w_p2_bar = torch.clamp(p2_bar, max=0)

    # G 0-(1,1), 1-(1,2), 2-(2,1), 3-(2,2)
    lambda1 = torch.where(G[2]*w_p2_bar < G[3]*p1_bar, torch.zeros_like(p1_bar), torch.where(G[1]*w_p1_bar < G[0]*p2_bar, w_p1_bar/G[0], torch.clamp(G[3]*p1_bar - G[2]*p2_bar, max=0)/(G[0]*G[3] - G[1]*G[2])))
    
    lambda2 = torch.where(G[2]*w_p2_bar < G[3]*p1_bar, w_p2_bar/G[3], torch.where(G[1]*w_p1_bar < G[0]*p2_bar, torch.zeros_like(p1_bar), torch.clamp(G[0]*p2_bar - G[1]*p1_bar, max=0)/(G[0]*G[3] - G[1]*G[2])))

    out = lambda1*y1_bar + lambda2*y2_bar + u_bar
    rt = xp1.clone()    
    rt[:,2:4] = x[:,2:4] + out
    # print(out)
    rt = rt.unsqueeze(0)
    return rt        

@torch.no_grad()
def p_sample(self, x, cond, t):
    b, *_, device = *x.shape, x.device
    model_mean, _, model_log_variance = self.p_mean_variance(x=x, cond=cond, t=t)
    noise = torch.randn_like(x)
    # no noise when t == 0
    nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))

    xp1 = model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise

    # Note:  choose any one of the below
    #---------------------------------------start--------------------------------------------------#
    ####################### original diffuser only
    # x = xp1      
    # xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    # yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    # off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    # off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1
    # b = ((x[:,2:3] - off_y)/yr)**2 + ((x[:,3:4] - off_x)/xr)**2 - 1
    # self.safe1 = torch.min(b[:,0])
    # xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    # yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    # off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    # off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1
    # b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1
    # self.safe2 = torch.min(b[:,0])

    ####################### truncate (shield) and GD (classifier-guidance/potential-based)
    # x = self.Shield(x, xp1)
    # x = self.GD(x, xp1)

    ####################### SafeDiffusers 
    #x = xp1 # for training only
    # x = self.invariance(x, xp1)    # RoS
    # x = self.invariance_cf(x, xp1)  # RoS closed form
    # x = self.invariance_relax(x, xp1, t) # ReS
    # x = self.invariance_relax_cf(x, xp1, t)   #ReS closed form    
    # x = self.invariance_time(x, xp1, t)   # TVS
    # x = self.invariance_time_cf(x, xp1, t)  # TVS closed form
    # x = self.invariance_relax_narrow(x, xp1, t)  # narrow passage case

    ####################### Applying SafeDiffusers to only the last 10 steps
    # if t <= 10:  #10
    #     # x = self.invariance_relax(x, xp1, t)  #done
    #     # x = self.invariance_relax_narrow(x, xp1, t)

    #     x = self.GD(x, xp1)
    # else:
    #     x = xp1
    #     xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    #     yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    #     off_x = 2*(5.8-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    #     off_y = 2*(5-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1
    #     b = ((x[:,2:3] - off_y)/yr)**2 + ((x[:,3:4] - off_x)/xr)**2 - 1
    #     self.safe1 = torch.min(b[:,0])
    #     xr = 2*1/(self.norm_maxs[1] - self.norm_mins[1])
    #     yr = 2*1/(self.norm_maxs[0] - self.norm_mins[0])
    #     off_x = 2*(5.3-0.5 - self.norm_mins[1])/(self.norm_maxs[1] - self.norm_mins[1]) - 1
    #     off_y = 2*(2-0.5 - self.norm_mins[0])/(self.norm_maxs[0] - self.norm_mins[0]) - 1
    #     b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1
    #     self.safe2 = torch.min(b[:,0])

    
    ###################### umaze case
    #x = self.invariance_umaze(x, xp1)   #umaze
    #x = self.GD(x, xp1)
    #x = self.Shield(x, xp1)  # Truncate method
    x = self.invariance_umaze_relax(x, xp1, t)   #umaze
    #-----------------------------------------end--------------------------------------------------#
    return x

@torch.no_grad()
def p_sample_loop(self, shape, cond, verbose=True, return_diffusion=False):
    device = self.betas.device

    batch_size = shape[0]
    x = torch.randn(shape, device=device)
    x = apply_conditioning(x, cond, self.action_dim)

    if return_diffusion: diffusion = [x]

    progress = utils.Progress(self.n_timesteps) if verbose else utils.Silent()
    safe1, safe2 = [], []
    for i in reversed(range(0, self.n_timesteps)):  #-50 change here for the number of diffusion steps,
        if i < 0:
            i = 0
        timesteps = torch.full((batch_size,), i, device=device, dtype=torch.long)
        x = self.p_sample(x, cond, timesteps)
        x = apply_conditioning(x, cond, self.action_dim)
        safe1.append(self.safe1.unsqueeze(0))
        safe2.append(self.safe2.unsqueeze(0))
        progress.update({'t': i})

        if return_diffusion: diffusion.append(x)
    
    self.safe1 = torch.cat(safe1, dim=0)
    self.safe2 = torch.cat(safe2, dim=0)

    progress.close()
    # pdb.set_trace()
    if return_diffusion:
        return x, torch.stack(diffusion, dim=1)
    else:
        return x

@torch.no_grad()
def conditional_sample(self, cond, *args, horizon=None, return_diffusion = True, **kwargs):
    '''
        conditions : [ (time, state), ... ]
    '''
    device = self.betas.device
    batch_size = len(cond[0])
    horizon = horizon or self.horizon
    shape = (batch_size, horizon, self.transition_dim)

    return self.p_sample_loop(shape, cond, return_diffusion= return_diffusion, *args, **kwargs)   ## debug

#------------------------------------------ training ------------------------------------------#

def q_sample(self, x_start, t, noise=None):
    if noise is None:
        noise = torch.randn_like(x_start)

    sample = (
        extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
        extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
    )

    return sample

def p_losses(self, x_start, cond, t):
    noise = torch.randn_like(x_start)

    x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
    x_noisy = apply_conditioning(x_noisy, cond, self.action_dim)

    x_recon = self.model(x_noisy, cond, t)
    x_recon = apply_conditioning(x_recon, cond, self.action_dim)

    assert noise.shape == x_recon.shape

    if self.predict_epsilon:
        loss, info = self.loss_fn(x_recon, noise)
    else:
        loss, info = self.loss_fn(x_recon, x_start)

    return loss, info

def loss(self, x, cond):
    batch_size = len(x)
    t = torch.randint(0, self.n_timesteps, (batch_size,), device=x.device).long()
    return self.p_losses(x, cond, t)


def _vb_terms_bpd(
    self, x_start, conditions, x_t, t, clip_denoised=True, model_kwargs=None
):
    """
    Get a term for the variational lower-bound.

    The resulting units are bits (rather than nats, as one might expect).
    This allows for comparison to other papers.

    :return: a dict with the following keys:
             - 'output': a shape [N] tensor of NLLs or KLs.
             - 'pred_xstart': the x_0 predictions.
    """
    # batch_size = x_start.shape(0)
    # conditions = self._format_conditions(conditions, batch_size)

    true_mean, _, true_log_variance_clipped = self.q_posterior(
        x_start=x_start, x_t=x_t, t=t
    )
    
    mean, _, log_variance = self.p_mean_variance(
         x_t, conditions, t)
    kl = normal_kl(
        true_mean, true_log_variance_clipped, mean, log_variance
    )
    kl = mean_flat(kl) / np.log(2.0)

    # import pdb; pdb.set_trace()
    # decoder_nll = -discretized_gaussian_log_likelihood(
    #     x_start, means=mean, log_scales=0.5 * log_variance
    # )
    
    # assert decoder_nll.shape == x_start.shape
    # decoder_nll = mean_flat(decoder_nll) / np.log(2.0)

    # At the first timestep return the decoder NLL,
    # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))

    # output = torch.where((t == 0), decoder_nll, kl)

    return kl


def forward(self, cond, *args, **kwargs):
    return self.conditional_sample(cond=cond, *args, **kwargs)

[Weixy21]

This involves another method which is hard to explain in plain words. I would suggest you to first read this paper (https://ieeexplore.ieee.org/abstract/document/9516971) and then you will know how to change it.

[userxzx2004]

thank you for your suggestion

[userxzx2004]

I've skimmed through the paper you provided, and I have a question: I'm using the official standard map and the default obstacle information. Do I need to adjust the parameters myself?

[Weixy21]

You first need to understand what is a CBF, then you will know how to change it.

[userxzx2004]

Thank you for your reply.

[userxzx2004]

Thank you again for your previous answer, but I still encountered a problem when I tried it.I tried this method，x = self.invariance_umaze_relax(x, xp1, t) #umaze，(diff) abc@abc:~/SafeDiffuser$ python -W ignore::UserWarning scripts/plan_maze2d.py --config config.maze2d --dataset maze2d-umaze-v1 --logbase logs
pybullet build time: Jan 29 2025 23:19:57
[ utils/colab ] Warning: not importing colab dependencies
[ utils/setup ] Reading config: config.maze2d:maze2d_umaze_v1
[ utils/setup ] Using overrides | config: config.maze2d | dataset: maze2d_umaze_v1
[ utils/setup ] Found extras: ['--logbase', 'logs']
[ utils/setup ] Overriding config | logbase : logs --> logs
[ utils/setup ] Lazy fstring | diffusion_loadpath : f:diffusion/H{horizon}_T{n_diffusion_steps} --> diffusion/H128_T64
[ utils/setup ] Setting exp_name to: plans/release_H128_T64_LimitsNormalizer_b1_condFalse
[ utils/setup ] Saved args to logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/args.json
[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/dataset_config.pkl

[utils/config ] Config: <class 'diffuser.datasets.sequence.GoalDataset'>
env: maze2d-umaze-v1
horizon: 128
max_path_length: 40000
normalizer: LimitsNormalizer
preprocess_fns: ['maze2d_set_terminals']
use_padding: False

[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/render_config.pkl

[utils/config ] Config: <class 'diffuser.utils.rendering.Maze2dRenderer'>
env: maze2d-umaze-v1

[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/model_config.pkl

[utils/config ] Config: <class 'diffuser.models.temporal.TemporalUnet'>
cond_dim: 4
dim_mults: (1, 4, 8)
horizon: 128
transition_dim: 6

[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/diffusion_config.pkl

[utils/config ] Config: <class 'diffuser.models.diffusion.GaussianDiffusion'>
action_dim: 2
action_weight: 1
clip_denoised: True
horizon: 128
loss_discount: 1
loss_type: l2
loss_weights: None
n_timesteps: 64
observation_dim: 4
predict_epsilon: False

[ utils/serialization ] Loaded config from logs/maze2d-umaze-v1/diffusion/H128_T64/trainer_config.pkl

[utils/config ] Config: <class 'diffuser.utils.training.Trainer'>
bucket: None
ema_decay: 0.995
gradient_accumulate_every: 2
label_freq: 40000
n_reference: 50
n_samples: 10
results_folder: logs/maze2d-umaze-v1/diffusion/H128_T64
sample_freq: 0
save_freq: 1000
save_parallel: False
train_batch_size: 32
train_lr: 0.0002

load datafile: 100%|████████████████████████████████████████████████████████| 8/8 [00:00<00:00, 31.26it/s]
[ utils/preprocessing ] Segmented maze2d-umaze-v1 | 1565 paths | min length: 96 | max length: 3993
[ datasets/buffer ] Finalized replay buffer | 1566 episodes
[ datasets/buffer ] Fields:
actions: (1566, 40000, 2)
infos/goal: (1566, 40000, 2)
infos/qpos: (1566, 40000, 2)
infos/qvel: (1566, 40000, 2)
observations: (1566, 40000, 4)
rewards: (1566, 40000, 1)
terminals: (1566, 40000, 1)
timeouts: (1566, 40000, 1)
next_observations: (1566, 40000, 4)
normed_observations: (1566, 40000, 4)
normed_actions: (1566, 40000, 2)
[ models/temporal ] Channel dimensions: [(6, 32), (32, 128), (128, 256)]
[(6, 32), (32, 128), (128, 256)]

[ utils/serialization ] Loading model epoch: 1960000

step: 0 /100
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/0.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/0.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/1.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/2.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/3.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/4.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/5.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/6.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/7.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/8.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/9.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/10.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/11.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/12.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/13.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/14.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/15.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/16.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/17.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/18.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/19.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/20.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/21.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/22.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/23.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/24.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/25.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/26.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/27.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/28.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/29.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/30.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/31.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/32.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/33.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/34.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/35.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/36.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/37.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/38.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/39.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/40.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/41.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/42.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/43.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/44.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/45.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/46.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/47.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/48.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/49.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/50.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/51.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/52.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/53.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/54.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/55.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/56.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/57.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/58.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/59.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/60.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/61.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/62.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/63.png
Saved 1 samples to: logs/maze2d-umaze-v1/plans/release_H128_T64_LimitsNormalizer_b1_condFalse/58/png/64.png
elbo mean: 0.0
elbo std: 0.0
safe1: 0.48210964
safe2: 0.67374897
score mean: 1.2256358234910512
score std: 0.0
computation time: 2.85821270942688
success rate: 1
(diff) abc@abc:~/SafeDiffuser$

The result will still be that it passes through the red box.

I tried modifying these parts of the code. if t >= 10: sign = 100 #relax b = ((x[:,2:3] - off_y)/yr)**4 + ((x[:,3:4] - off_x)/xr)**4 - 1 It still hasn't achieved the effect the barrier function should have. I would be very grateful if you could provide some more specific suggestions.

[userxzx2004]

If I use this method directly，x = self.invariance_umaze(x, xp1)，It seems to meet the requirements, but the generated trajectory goes through the walls.

[userxzx2004]

could you give me some suggestions