PyTorch implementation for our paper "Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization"
import copy
from time import time
from hydra.utils import instantiate

import numpy as np
import matplotlib.pyplot as plt
import torch
from torch import nn
from einops import rearrange, reduce, repeat

from gflownet.network import FourierMLP, TimeConder
from gflownet.utils import normal_logp, num_to_groups, loss2ess_info
from target.plot import viz_sample2d, viz_kde2d, viz_contour_sample2d


def get_alg(cfg, task=None):
    alg = SubTrajectoryBalanceTransitionBased(cfg, task=task)
    return alg

def fl_inter_logr(x, logreward_fn, config, cur_t, sigma=None): # x: (bs, dim)
    if sigma is None:
        sigma = config.sigma

    if isinstance(cur_t, torch.Tensor):
        # print(cur_t.shape) # (bs, 1) or (,)
        if cur_t.ndim <= 1:
            cur_t = cur_t.item()
    ratio = cur_t / config.t_end
    # assert 0 <= ratio <= 1

    coef = max(np.sqrt(0.01 * config.t_end), np.sqrt(cur_t))  # cur_t could be 0
    logp0 = normal_logp(x, 0., coef * sigma)
    fl_logr = logreward_fn(x) * ratio + logp0 * (1 - ratio)

    return fl_logr

def sample_traj(gfn, config, logreward_fn, batch_size=None, sigma=None):
    if batch_size is None:
        batch_size = config.batch_size
    device = gfn.device
    if sigma is None:
        sigma = config.sigma

    x = gfn.zero(batch_size).to(device)
    fl_logr = fl_inter_logr(x, logreward_fn, config, cur_t=0., sigma=sigma)
    traj = [(torch.tensor(0.), x.cpu(), fl_logr.cpu())] # (t, x, logr)
    inter_loss = torch.zeros(batch_size).to(device)

    x_max = 0.
    for cur_t in torch.arange(0, config.t_end, config.dt).to(device):
        x, uw_term, u2_term = gfn.step_forward(cur_t, x, config.dt, sigma=sigma)
        x = x.detach()
        fl_logr = fl_inter_logr(x, logreward_fn, config,
            cur_t=cur_t + config.dt, sigma=sigma).detach().cpu()
        traj.append((cur_t.cpu() + config.dt, x.detach().cpu(), fl_logr))
        inter_loss += (u2_term + uw_term).detach()
        x_max = max(x_max, x.abs().max().item())

    pis_terminal = -gfn.nll_prior(x) - logreward_fn(x)
    pis_log_weight = inter_loss + pis_terminal
    info = {"pis_logw": pis_log_weight, "x_max": x_max}
    return traj, info


class GFlowNet(nn.Module):
    """
    For PIS modeling: s0 is fixed to be zero
    thus PB(s0|s1) == 1
    """
    def __init__(self, cfg, task=None):
        super().__init__()
        self.cfg = cfg

        self.data_ndim = cfg.data_ndim  # int(np.prod(data_shape))
        self.register_buffer("x0", torch.zeros((1, self.data_ndim))) # for pis modeling
        self.t_end = cfg.t_end
        self.dt = cfg.dt

        self.g_func = instantiate(cfg.g_func)
        self.f_func = instantiate(cfg.f_func) # learnable
        self.nn_clip = cfg.nn_clip
        self.lgv_clip = cfg.lgv_clip
        self.task = task
        self.grad_fn = task.score
        self.select_f(f_format=cfg.f)
        self.xclip = cfg.xclip
        self.logr_fn = lambda x: -task.energy(x)

    @property
    def device(self):
        return next(self.parameters()).device

    @property
    def sigma(self): # simply return a float w/o grad
        return self.cfg.sigma

    def save(self, path):
        raise NotImplementedError

    def get_optimizer(self):
        return torch.optim.Adam(self.param_ls, weight_decay=self.cfg.weight_decay)

    def select_f(self, f_format=None):
        _fn = self.f_func
        if f_format == "f":
            def _fn(t, x):
                return torch.clip(self.f_func(t, x), -self.nn_clip, self.nn_clip)
        elif f_format in ["t_tnet_grad", "tgrad"]:
            self.lgv_coef = TimeConder(self.cfg.f_func.channels, 1, 3)
            def _fn(t, x):
                grad = torch.nan_to_num(self.grad_fn(x))
                grad = torch.clip(grad, -self.lgv_clip, self.lgv_clip)

                f = torch.clip(self.f_func(t, x), -self.nn_clip, self.nn_clip)
                lgv_coef = self.lgv_coef(t)
                return f - lgv_coef * grad
        else:
            raise RuntimeError
        self.param_fn = _fn

    # t: scalar; state: (b, state_dim)
    def f(self, t, state): # same as self.param_fn
        x = torch.nan_to_num(state)
        control = self.param_fn(t, x) #.flatten(start_dim=1)
        return control

    def zero(self, batch_size, device=None):
        device = self.device if device is None else device
        return self.x0.expand(batch_size, -1).to(device)

    def nll_prior(self, state): # nll on terminal state
        return -normal_logp(state, 0., np.sqrt(self.t_end) * self.sigma)

    # t -> t + dt; n -> n + 1; here t is scalar tensor
    def step_forward(self, t, state, dt, sigma=None, return_drift_scale=False):
        sigma = self.sigma if sigma is None else sigma
        std_noise = self.g_func(t, state) * torch.randn_like(state)

        noise_term = std_noise * sigma * np.sqrt(dt)
        pre_drift = self.f(t, state)
        # drift_scale = pre_drift.norm(dim=-1).mean()
        next_state = state + pre_drift * sigma * dt + noise_term
        u2_term = 0.5 * pre_drift.pow(2).sum(dim=-1) * dt
        uw_term = (pre_drift * std_noise).sum(dim=1) * np.sqrt(dt)

        if self.xclip > 0: # avoid nan
            next_state = torch.clip(next_state, -self.xclip, self.xclip)

        return next_state, uw_term, u2_term

    # t -> t - dt; n -> n - 1
    def step_backward(self, t ,state, dt): # not used
        std_noise = self.g_func(t, state) * torch.randn_like(state)
        # n = (t - dt) / dt
        coef = (t-dt)/t  # = n/(n + 1)
        mean = self.x0*dt / t + coef * state
        noise_term = std_noise * self.sigma * np.sqrt(dt)
        prev_state = mean + coef.sqrt() * noise_term
        return prev_state.detach()

    def log_pf(self, t, state, next_state, dt=None): # t: (bs, 1), dt: float
        assert t.ndim == state.ndim == next_state.ndim == 2
        dt = self.dt if dt is None else dt
        sigma = self.cfg.sigma
        mean = state + self.f(t, state) * sigma * dt
        log_pf = normal_logp(next_state, mean, sigma * np.sqrt(dt))
        return log_pf

    def log_pb(self, t, state, prev_state, dt=None): # t: (bs, 1), dt: float
        assert t.ndim == state.ndim == prev_state.ndim == 2
        dt = self.dt if dt is None else dt
        sigma = self.cfg.sigma

        mean = self.x0 * dt / t + (t - dt) / t * state
        sigma_pb = ((t-dt)/t).sqrt() * np.sqrt(dt) * sigma
        log_pb = normal_logp(prev_state, mean, sigma_pb)

        # first step is from Dirac on x0 to a Gaussian, thus PB == 1
        # sigma_pb.min = 0 => nan
        first_step_mask = (t <= dt).squeeze(dim=1)
        # log_pb[first_step_mask] = 0.
        log_pb = torch.where(first_step_mask, torch.zeros_like(log_pb), log_pb)
        return log_pb

    def log_pf_and_pb_traj(self, traj):
        batch_size = traj[0][1].shape[0]

        xs = [x for (t, x, r) in traj]
        state = torch.cat(xs[:-1], dim=0).to(self.device)  # (N*b, d)
        next_state = torch.cat(xs[1:], dim=0).to(self.device)  # (N*b, d)
        ts = [repeat(t[None], "one -> b one", b=batch_size) for (t, x, r) in traj]
        time = torch.cat(ts[:-1], dim=0).to(self.device)  # (N*b, 1)
        next_time = torch.cat(ts[1:], dim=0).to(self.device)  # (N*b, 1)

        log_pb = self.log_pb(next_time, next_state, state)
        log_pb = log_pb.detach()

        if self.cfg.task in ["cox"]:  # save cuda memory
            start_idx = 0
            log_pf = torch.zeros((0,)).to(self.device)
            for bs in num_to_groups(log_pb.shape[0], 1000):
                log_pf_curr = self.log_pf(time[start_idx:start_idx+bs],
                          state[start_idx:start_idx+bs], next_state[start_idx:start_idx+bs])
                log_pf = torch.cat([log_pf, log_pf_curr], dim=0)
                start_idx += bs
        else:
            log_pf = self.log_pf(time, state, next_state)

        return log_pf, log_pb

    def log_weight(self, traj): # "log q - log p" in VAE notation
        batch_size = traj[0][1].shape[0]
        logr = self.logr_from_traj(traj)
        log_pf, log_pb = self.log_pf_and_pb_traj(traj) # (N*b,)
        dlogp = reduce(log_pf - log_pb, "(N b) -> b", "sum", b=batch_size)
        return dlogp - logr

    def logr_from_traj(self, traj):
        return traj[-1][2].to(self.device)

    @torch.no_grad()
    def eval_step(self, num_samples, logreward_fn=None):
        if logreward_fn is None:
            logreward_fn = self.logr_fn
        self.eval()

        if self.cfg.task in ["cox"]: # save cuda memory
            bs_ls = num_to_groups(num_samples, 250)
            pis_logw = None
            x_max = 0.
            for bs in bs_ls:
                traj_curr, sample_info = sample_traj(self, self.cfg, logreward_fn, batch_size=bs)
                logw_curr = sample_info['pis_logw']
                x_max = max(x_max, sample_info['x_max'])
                if pis_logw is None:
                    pis_logw = logw_curr
                    traj = traj_curr
                else:
                    pis_logw = torch.cat([pis_logw, logw_curr], dim=0)

            print(f"logw_pis={pis_logw.mean().item():.8e}")
            logw = pis_logw
        else:
            traj, sample_info = sample_traj(self, self.cfg, logreward_fn, batch_size=num_samples)
            pis_logw = sample_info['pis_logw']
            x_max = sample_info['x_max']
            logw = self.log_weight(traj)
            print(f"logw={logw.mean().item():.8e}, logw_pis={pis_logw.mean().item():.8e}")

        # pis_logZ = torch.logsumexp(-pis_logw, dim=0) - np.log(num_samples)
        # Z = \int R(x) dx = E_{PF(tau)}[R(x)PB(tau|x)/PF(tau)]]
        logZ_eval = torch.logsumexp(-logw, dim=0) - np.log(num_samples) # (bs,) -> ()
        logZ_elbo = torch.mean(-logw, dim=0)
        info = {"logz": logZ_eval.item(),
                "logz_elbo": logZ_elbo.item(),
                "x_max": x_max}

        return traj, info

    def visualize(self, traj, logreward_fn=None, step=-1):
        state = traj[-1][1].detach().cpu()
        step_str = (f"-{step}" if step >= 0 else "")

        data_ndim = self.data_ndim
        lim = 7
        if self.cfg.task in ["funnel"]:
            data_ndim = 2
            state = state[:, :2]

        if data_ndim == 2:
            dist_img_path = f"dist{step_str}.png"
            viz_sample2d(state, None, dist_img_path, lim=lim)
            viz_sample2d(state, None, f"dist{step_str}.pdf", lim=lim)

            kde_img_path = f"kde{step_str}.png"
            viz_kde2d(state, None, kde_img_path, lim=lim)
            viz_kde2d(state, None, f"kde{step_str}.pdf", lim=lim)

            alpha = 0.8
            n_contour_levels = 20
            def logp_func(x): return -self.task.energy(x.cuda()).cpu()
            contour_img_path = f"contour{step_str}.png"
            viz_contour_sample2d(state, contour_img_path, logp_func, lim=lim, alpha=alpha,
                                 n_contour_levels=n_contour_levels)
            viz_contour_sample2d(state, f"contour{step_str}.pdf", logp_func, lim=lim, alpha=alpha,
                                 n_contour_levels=n_contour_levels)

            img_dict = {"distribution": dist_img_path,
                        "KDE": kde_img_path,
                        "contour":contour_img_path,
                        }

            return img_dict

        if self.cfg.task in ["manywell"]:
            lim = 3
            x13 = state[:, 0:3:2] # 1st and 3rd dimension
            dist_img_path = f"distx13{step_str}.png"
            viz_sample2d(x13, None, dist_img_path, lim=lim)
            kde_img_path = f"kdex13{step_str}.png"
            viz_kde2d(x13, None, kde_img_path, lim=lim)

            alpha = 0.8
            n_contour_levels = 20
            def logp_func(x_2d):
                x = torch.zeros((x_2d.shape[0], self.data_ndim))
                x[:, 0] = x_2d[:, 0]
                x[:, 2] = x_2d[:, 1]
                return -self.task.energy(x)
            contour_img_path = f"contourx13{step_str}.png"
            viz_contour_sample2d(x13, contour_img_path, logp_func, lim=lim, alpha=alpha, n_contour_levels=n_contour_levels)
            viz_contour_sample2d(x13, f"contourx13{step_str}.pdf", logp_func, lim=lim, alpha=alpha, n_contour_levels=n_contour_levels)

            x23 = state[:, 1:3] # 2nd and 3rd dimension
            dist_img_path2 = f"distx23{step_str}.png"
            viz_sample2d(x23, None, dist_img_path2, lim=lim)
            viz_kde2d(x23, None, f"kdex23{step_str}.png", lim=lim)

            def logp_func(x_2d):
                x = torch.zeros((x_2d.shape[0], self.data_ndim))
                x[:, 1] = x_2d[:, 0]
                x[:, 2] = x_2d[:, 1]
                return -self.task.energy(x)
            contour_img_path2 = f"contourx23{step_str}.png"
            viz_contour_sample2d(x23, contour_img_path2, logp_func, lim=lim, alpha=alpha, n_contour_levels=n_contour_levels)
            viz_contour_sample2d(x23, f"contourx23{step_str}.pdf", logp_func, lim=lim, alpha=alpha, n_contour_levels=n_contour_levels)

            return {"distribution": dist_img_path,
                    "distribution2": dist_img_path2,
                    "KDE": kde_img_path,
                    "contour": contour_img_path,
                    "contour2": contour_img_path2,
                    }


class DetailedBalance(GFlowNet):
    def __init__(self, cfg, task=None):
        super().__init__(cfg, task)
        self.flow = FourierMLP(
            self.data_ndim, 1,
           num_layers=cfg.f_func.num_layers, channels=cfg.f_func.channels,
           zero_init=True
        )
        self.param_ls = [
            {"params": self.f_func.parameters(), "lr": self.cfg.lr},
            {"params": self.flow.parameters(), "lr": self.cfg.zlr},
        ]
        if hasattr(self, "lgv_coef"):
            self.param_ls.append({"params": self.lgv_coef.parameters(), "lr": self.cfg.lr})
        self.optimizer = self.get_optimizer()

    def save(self, path="alg.pt"):
        self.eval()
        save_dict = {
            "f_func": self.f_func.state_dict(),
            "flow": self.flow.state_dict(),
        }
        torch.save(save_dict, path)

    def load(self, path="alg.pt"):
        save_dict = torch.load(path)
        self.f_func.load_state_dict(save_dict["f_func"])
        self.flow.load_state_dict(save_dict["flow"])



def cal_subtb_coef_matrix(lamda, N):
    """
    diff_matrix: (N+1, N+1)
    0, 1, 2, ...
    -1, 0, 1, ...
    -2, -1, 0, ...

    self.coef[i, j] = lamda^(j-i) / total_lambda  if i < j else 0.
    """
    range_vals = torch.arange(N + 1)
    diff_matrix = range_vals - range_vals.view(-1, 1)
    B = np.log(lamda) * diff_matrix
    B[diff_matrix <= 0] = -np.inf
    log_total_lambda = torch.logsumexp(B.view(-1), dim=0)
    coef = torch.exp(B - log_total_lambda)
    return coef


class SubTrajectoryBalanceTransitionBased(DetailedBalance):
    def __init__(self, cfg, task=None):
        super().__init__(cfg, task)

        self.Lambda = float(cfg.subtb_lambda)
        coef = cal_subtb_coef_matrix(self.Lambda, int(cfg.N)) # (N+1, N+1)
        self.register_buffer('coef', coef, persistent=False)

    def train_step(self, traj, debug=False):
        self.train()
        batch_size = traj[0][1].shape[0]

        xs = [x.to(self.device) for (t, x, r) in traj]
        states = torch.cat(xs, dim=0)  # ((N+1)*b, d)
        states = rearrange(states, "(T1 b) d -> b T1 d", b=batch_size)  # (b, N+1, d)

        ts = [t[None].to(self.device) for (t, x, r) in traj]
        times = torch.cat(ts, dim=0)  # (N+1, 1)
        time_b = repeat(times, "T1 -> T1 one", one=1)  # (N+1, 1)
        # time_coef = (time_b / self.t_end).squeeze(-1)  # (N+1, 1) -> (N+1,)

        rs = [r.to(self.device) for (t, x, r) in traj]
        logrs = torch.cat(rs, dim=0)  # ((N+1)*b, 1)
        logrs = rearrange(logrs, "(T1 b) -> b T1", b=batch_size)  # (b, N+1)

        info = {"logR": logrs[:, -1].mean().item()}
        for b_idx in range(batch_size):
            state_b = states[b_idx]  # (N+1, d)
            log_pf = self.log_pf(time_b[:-1], state_b[:-1], state_b[1:])  # (N,)
            log_pb = self.log_pb(time_b[1:], state_b[1:], state_b[:-1])  # (N,)

            flow_b = self.flow(time_b, state_b).squeeze(-1)  # (N+1, 1) -> (N+1,)

            flow_b = flow_b + logrs[b_idx]  # (N+1,)
            flow_b[-1] = logrs[b_idx][-1]  # (1,)

            diff_logp = log_pf - log_pb # (N, )
            diff_logp_padded = torch.cat(
                (torch.zeros(1).to(diff_logp), diff_logp.cumsum(dim=-1))
            , dim=0) # (N+1,)
            # this means A1[i, j] = diff_logp[i:j].sum(dim=-1)
            A1 = diff_logp_padded.unsqueeze(0) - diff_logp_padded.unsqueeze(1)  # (N+1, N+1)

            A2 = flow_b[:, None] - flow_b[None, :] + A1  # (N+1, N+1)
            if not torch.all(torch.isfinite(A2)):
                import ipdb;
                ipdb.set_trace()

            A2 = (A2 / self.data_ndim).pow(2)  # (N+1, N+1)
            # torch.triu() is useless here
            loss = torch.triu(A2 * self.coef, diagonal=1).sum()
            info["loss_dlogp"] = loss.item()

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

        info["loss_train"] = loss.item()
        return info
gflownet/gflownet.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
import sys, os
import gzip, pickle
from collections import defaultdict
import hydra
from hydra.utils import instantiate
from omegaconf import DictConfig, open_dict, OmegaConf

import numpy as np
import torch

from gflownet.gflownet import get_alg, sample_traj
from gflownet.utils import seed_torch

torch.backends.cudnn.benchmark = True


def refine_cfg(cfg):
    with (open_dict(cfg)):
        cfg.device = cfg.d
        cfg.work_directory = os.getcwd()
        cfg.gpu_type = torch.cuda.get_device_name()\
            if (torch.cuda.is_available() and cfg.device >= 0) else "CPU"
        print(f"GPU type: {cfg.gpu_type}")

        cfg.task = cfg.target.dataset._target_.split(".")[-2]
        cfg.logr_min = cfg.rmin

        cfg.batch_size = cfg.bs
        cfg.weight_decay = cfg.wd
        cfg.sigma_interactive = cfg.sgmit
        if cfg.sigma_interactive <= 0:
            cfg.sigma_interactive = cfg.sigma
        cfg.t_end = cfg.dt * cfg.N
        cfg.subtb_lambda = cfg.stlam

    del cfg.d, cfg.bs, cfg.wd, cfg.sgmit, cfg.stlam, cfg.rmin
    return cfg


@hydra.main(config_path="configs", config_name="main")
def main(cfg: DictConfig) -> None:
    cfg = refine_cfg(cfg)
    device = torch.device(f"cuda:{cfg.device:d}" if torch.cuda.is_available() and cfg.device>=0 else "cpu")

    seed_torch(cfg.seed)
    print(f"Device: {device}, GPU type: {cfg.gpu_type}")
    print(str(cfg))
    print(f"Work directory: {os.getcwd()}")

    data = instantiate(cfg.target.dataset)
    true_logz = data.gt_logz()
    if true_logz is not None:
        print(f"True logZ={true_logz:.4f}")
    def logr_fn_detach(x):
        logr = -data.energy(x).detach()
        logr = torch.where(torch.isinf(logr), cfg.logr_min, logr)
        logr = torch.clamp(logr, min=cfg.logr_min)
        return logr
    gflownet = get_alg(cfg, data)
    gflownet.to(device)

    metric_best = 100.
    for step_idx in range(cfg.steps):
        ######### eval
        gflownet.eval()
        if step_idx % cfg.eval_freq == 0 or step_idx >= cfg.steps-1:
            traj, eval_info = gflownet.eval_step(cfg.eval_n, logr_fn_detach)
            print(f"EVALUATION: step={step_idx}:",
                  " ".join([f"{k}={v:.3f}" for k, v in eval_info.items()]))

            if true_logz is not None:
                logz_diff = abs(eval_info["logz"] - true_logz)
                print(f"logZ diff={logz_diff:.3f}")
                if logz_diff < metric_best:
                    metric_best = logz_diff
                    print(f"best metric: {metric_best:.2f} at step {step_idx}")

        ######### rollout
        traj_batch, _ = sample_traj(gflownet, cfg, logr_fn_detach,
              batch_size=cfg.batch_size, sigma = cfg.sigma_interactive)

        ######### training
        gflownet.train()
        gflownet.train_step(traj_batch)


if __name__ == "__main__":
    main()
gflownet/main.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
from collections.abc import Iterable

import numpy as np
import torch
from torch import nn
from einops import rearrange


def check_shape(cur_shape):
    if isinstance(cur_shape, Iterable):
        return tuple(cur_shape)
    elif isinstance(cur_shape, int):
        return tuple([cur_shape,])
    else:
        raise NotImplementedError(f"Type {type(cur_shape)} not support")

class IdentityOne:
    def __call__(self, t, y):
        del t
        return torch.ones_like(y)

class TimeConder(nn.Module):
    def __init__(self, channel, out_dim, num_layers):
        super().__init__()
        self.register_buffer(
            "timestep_coeff", torch.linspace(start=0.1, end=100, steps=channel)[None]
        )
        self.timestep_phase = nn.Parameter(torch.randn(channel)[None])
        self.layers = nn.Sequential(
            nn.Linear(2 * channel, channel),
            *[
                nn.Sequential(
                    nn.GELU(),
                    nn.Linear(channel, channel),
                )
                for _ in range(num_layers - 1)
            ],
            nn.GELU(),
            nn.Linear(channel, out_dim)
        )

        self.layers[-1].weight.data.fill_(0.0)
        self.layers[-1].bias.data.fill_(0.01)

    def forward(self, t):
        sin_cond = torch.sin((self.timestep_coeff * t.float()) + self.timestep_phase)
        cos_cond = torch.cos((self.timestep_coeff * t.float()) + self.timestep_phase)
        cond = rearrange([sin_cond, cos_cond], "d b w -> b (d w)")
        return self.layers(cond)


class FourierMLP(nn.Module):
    def __init__(self, in_shape, out_shape, num_layers=2, channels=128,
                 zero_init=True, res=False):
        super().__init__()
        self.in_shape = check_shape(in_shape) # 2 -> (2,)
        self.out_shape = check_shape(out_shape)

        self.register_buffer(
            "timestep_coeff", torch.linspace(start=0.1, end=100, steps=channels)[None]
        )
        self.timestep_phase = nn.Parameter(torch.randn(channels)[None])
        self.input_embed = nn.Linear(int(np.prod(in_shape)), channels)
        self.timestep_embed = nn.Sequential(
            nn.Linear(2 * channels, channels),
            nn.GELU(),
            nn.Linear(channels, channels),
        )
        self.layers = nn.Sequential(
            nn.GELU(),
            *[
                nn.Sequential(nn.Linear(channels, channels), nn.GELU())
                for _ in range(num_layers)
            ],
        )
        self.final_layer = nn.Linear(channels, int(np.prod(self.out_shape)))
        if zero_init:
            self.final_layer.weight.data.fill_(0.0)
            self.final_layer.bias.data.fill_(0.0)

        self.residual = res

    # cond: (1,) or (1, 1) or (bs, 1); inputs: (bs, d)
    # output: (bs, d_out)
    def forward(self, cond, inputs):
        cond = cond.view(-1, 1).expand((inputs.shape[0], 1))
        sin_embed_cond = torch.sin(
            # (1, channels) * (bs, 1) + (1, channels)
            (self.timestep_coeff * cond.float()) + self.timestep_phase
        )
        cos_embed_cond = torch.cos(
            (self.timestep_coeff * cond.float()) + self.timestep_phase
        )
        embed_cond = self.timestep_embed(
            rearrange([sin_embed_cond, cos_embed_cond], "d b w -> b (d w)")
        ) # (bs, 2* channels) -> (bs, channels)
        embed_ins = self.input_embed(inputs.view(inputs.shape[0], -1)) # (bs, d) -> (bs, channels)

        input = embed_ins + embed_cond
        out = self.layers(input)
        if self.residual:
            out = out + input
        out = self.final_layer(out) # (bs, channels) -> (bs, d)
        return out.view(-1, *self.out_shape)
gflownet/network.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
import os, sys
from itertools import count
import pathlib
import functools
import socket

import math
import random
import scipy
import numpy as np
import torch


######### System Utils

def num_to_groups(num, divisor):
    groups = num // divisor
    remainder = num % divisor
    arr = [divisor] * groups
    if remainder > 0:
        arr.append(remainder)
    return arr

######### Pytorch Utils

import random
def seed_torch(seed, verbose=True):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)

    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    # torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.

    if verbose:
        print("==> Set seed to {:}".format(seed))

def normal_logp(x, mean=0., sigma=1.):
    # x: (bs, dim)
    # mean: scalar or (bs, dim)
    # sigma: scalar float or (bs, 1); assume all dimensions have the same sigma

    assert x.ndim == 2
    # dim = x.shape[-1]
    if isinstance(sigma, torch.Tensor):
        assert sigma.ndim == 2
        log_sigma = torch.log(sigma)
    else:
        log_sigma = np.log(sigma)

    # broadcast: sigma (bs, 1) + mean (bs, dim) -> (bs, dim)
    neg_logp = 0.5 * np.log(2 * np.pi) + log_sigma \
               + 0.5 * (x - mean) ** 2 / (sigma ** 2)
    return torch.sum(-neg_logp, dim=-1) # (bs,)

def loss2ess_info(loss):
    # ESS = (\sum w_i)^2 / \sum w_i^2
    # return ESS / N <= 1
    log_weight = -loss + loss.mean()
    log_numerator = 2 * torch.logsumexp(log_weight, dim=0)
    log_denominator = torch.logsumexp(2 * log_weight, dim=0)
    ess = torch.exp(log_numerator - log_denominator) / len(log_weight)
    return {"ess": ess.item()}
gflownet/utils.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
import abc
import torch
from torch.utils.data import Dataset

from target.ksd import KSD
from target.loss_helper import linear_intepolate_energy, nll_unit_gaussian


class BaseSet(abc.ABC, Dataset):
    def __init__(self, len_data=-2333):
        self.num_sample = len_data
        self.data = None
        self.data_ndim = None
        self.worker = KSD(self.score, beta=0.2)
        self._gt_ksd = None

    def gt_logz(self):
        raise NotImplementedError

    @abc.abstractmethod
    def energy(self, x):
        return

    def unnorm_pdf(self, x):
        return torch.exp(-self.energy(x))

    # hmt stands for hamiltonian
    def hmt_energy(self, x):
        dim = x.shape[-1]
        x, v = torch.split(x, dim // 2, dim=-1)
        neg_log_p_x = self.sample_energy_fn(x)
        neg_log_p_v = nll_unit_gaussian(v)
        return neg_log_p_x + neg_log_p_v

    @property
    def ndim(self):
        return self.data_ndim

    def sample(self, batch_size):
        del batch_size
        raise NotImplementedError

    def score(self, x):
        with torch.no_grad():
            copy_x = x.detach().clone()
            copy_x.requires_grad = True
            with torch.enable_grad():
                self.energy(copy_x).sum().backward()
                lgv_data = copy_x.grad.data
            return lgv_data

    def hmt_score(self, x):
        with torch.no_grad():
            copy_x = x.detach().clone()
            copy_x.requires_grad = True
            with torch.enable_grad():
                self.hmt_energy(copy_x).sum().backward()
                lgv_data = copy_x.grad.data
            return lgv_data

    def ksd(self, points):
        with torch.no_grad():
            cur_ksd = self.gt_ksd()
        return self.worker(points) - cur_ksd

    def gt_ksd(self):
        if self._gt_ksd is None:
            with torch.no_grad():
                self._gt_ksd = self.worker(
                    self.sample(5000).view(5000, -1), adjust_beta=True
                )
        return self._gt_ksd

target/distribution/base_set.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
import numpy as np
import torch
import torch as th
import torch.distributions as D

from .base_set import BaseSet


class FunnelSet(BaseSet):
    """
    x0 ~ N(0, 3^2), xi | x0 ~ N(0, exp(x0)), i = 1, ..., 9
    """
    def __init__(self, dim):
        super().__init__()
        self.data = th.ones(dim, dtype=float).cuda()
        self.data_ndim = dim

        self.dist_dominant = D.Normal(th.tensor([0.0]).cuda(), th.tensor([3.0]).cuda())
        self.mean_other = th.zeros(dim - 1).float().cuda()
        self.cov_eye = th.eye(dim - 1).float().cuda().view(1, dim - 1, dim - 1)

    # def cal_gt_big_z(self):
    #     return 1

    def gt_logz(self):
        return 0.

    def energy(self, x):
        return -self.funner_log_pdf(x)

    def viz_pdf(self, fsave="density.png", lim=3):
        raise NotImplementedError

    def funner_log_pdf(self, x):
        try:
            dominant_x = x[:, 0]
            log_density_dominant = self.dist_dominant.log_prob(dominant_x)  # (B, )
            # log_density_other = self._dist_other(dominant_x).log_prob(x[:, 1:])  # (B, )

            log_sigma = 0.5 * x[:, 0:1]
            sigma2 = torch.exp(x[:, 0:1])
            neglog_density_other = 0.5*np.log(2*np.pi) + log_sigma + 0.5 * x[:, 1:] ** 2 / sigma2
            log_density_other = torch.sum(-neglog_density_other, dim=-1)
        except:
            import ipdb;
            ipdb.set_trace()
        return log_density_dominant + log_density_other

    def sample(self, batch_size):
        dominant_x = self.dist_dominant.sample((batch_size,))  # (B,1)
        x_others = self._dist_other(dominant_x).sample()  # (B, dim-1)
        return th.hstack([dominant_x, x_others])

    def _dist_other(self, dominant_x):
        variance_other = th.exp(dominant_x)
        cov_other = variance_other.view(-1, 1, 1) * self.cov_eye
        # use covariance matrix, not std
        return D.multivariate_normal.MultivariateNormal(self.mean_other, cov_other)
target/distribution/funnel.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
import matplotlib.pyplot as plt

import torch
import torch.distributions as D
from torch.distributions.mixture_same_family import MixtureSameFamily

from .base_set import BaseSet

"""
2-D Guassian mixture
"""
class GaussianMixture2D(BaseSet):
    def __init__(self, scale=0.5477222, # sqrt(0.3)
                 # nmode=9, xlim=1.0
                 ):
        super().__init__()
        # xlim = 0.01 if nmode == 1 else xlim
        self.data = torch.tensor([0.0])

        mean_ls = [
            [-5., -5.], [-5., 0.], [-5., 5.],
            [0., -5.], [0., 0.], [0., 5.],
            [5., -5.], [5., 0.], [5., 5.],
        ]
        nmode = len(mean_ls)
        mean = torch.stack([torch.tensor(xy) for xy in mean_ls])
        comp = D.Independent(D.Normal(mean.cuda(), torch.ones_like(mean).cuda() * scale), 1)
        mix = D.Categorical(torch.ones(nmode).cuda())
        self.gmm = MixtureSameFamily(mix, comp)

    def gt_logz(self):
        return 0.

    def energy(self, x):
        return -self.gmm.log_prob(x).flatten()

    def sample(self, batch_size):
        return self.gmm.sample((batch_size,))

    def viz_pdf(self, fsave="ou-density.png"):
        x = torch.linspace(-8, 8, 100).cuda()
        y = torch.linspace(-8, 8, 100).cuda()
        X, Y = torch.meshgrid(x, y)
        x = torch.stack([X.flatten(), Y.flatten()], dim=1) #?

        density = self.unnorm_pdf(x)
        # x, pdf = as_numpy([x, density])
        x, pdf = torch.from_numpy(x), torch.from_numpy(density)

        fig, axs = plt.subplots(1, 1, figsize=(1 * 7, 1 * 7))
        axs.plot(x, pdf)

        # plt.contourf(X, Y, density, levels=20, cmap='viridis')
        # plt.colorbar()
        # plt.xlabel('x')
        # plt.ylabel('y')
        # plt.title('2D Function Plot')

        fig.savefig(fsave)
        plt.close(fig)

    def __getitem__(self, idx):
        del idx
        return self.data[0]

target/distribution/gm.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
from typing import Callable

import numpy as np
import torch
import torch.distributions as D

from .base_set import BaseSet
from target.plot import traj_plot1d, dist_plot1d, viz_sample2d, viz_kde2d, viz_contour_sample2d


def rejection_sampling(n_samples: int, proposal: torch.distributions.Distribution,
                       target_log_prob_fn: Callable, k: float) -> torch.Tensor:
    """Rejection sampling. See Pattern Recognition and ML by Bishop Chapter 11.1"""
    z_0 = proposal.sample((n_samples*10,))
    u_0 = torch.distributions.Uniform(0, k*torch.exp(proposal.log_prob(z_0)))\
        .sample().to(z_0)
    accept = torch.exp(target_log_prob_fn(z_0)) > u_0
    samples = z_0[accept]
    if samples.shape[0] >= n_samples:
        return samples[:n_samples]
    else:
        required_samples = n_samples - samples.shape[0]
        new_samples = rejection_sampling(required_samples, proposal, target_log_prob_fn, k)
        samples = torch.concat([samples, new_samples], dim=0)
        return samples

class ManyWell(BaseSet):
    """
    log p(x1, x2) = −x1^4 + 6*x1^2 + 1/2*x1 − 1/2*x2^2 + constant
    """
    def __init__(self, len_data, dim=32, is_linear=True):
        super().__init__(len_data)
        self.data = torch.ones(dim, dtype=float).cuda()
        self.data_ndim = dim
        assert dim % 2 == 0
        self.n_wells = dim // 2

        # as rejection sampling proposal
        self.component_mix = torch.tensor([0.2, 0.8])
        self.means = torch.tensor([-1.7, 1.7])
        self.scales = torch.tensor([0.5, 0.5])

        self.Z_x1 = 11784.50927
        self.logZ_x2 = 0.5*np.log(2*np.pi)
        self.logZ_doublewell = np.log(self.Z_x1) + self.logZ_x2

    def gt_logz(self):
        return self.n_wells * self.logZ_doublewell

    def energy(self, x):
        return -self.manywell_logprob(x)

    def doublewell_logprob(self, x):
        assert x.shape[1] == 2 and x.ndim == 2
        x1 = x[:, 0]
        x2 = x[:, 1]
        x1_term = 0.5*x1 + 6*x1.pow(2) - x1.pow(4)
        x2_term = -0.5*x2.pow(2)
        return x1_term + x2_term

    def manywell_logprob(self, x):
        assert x.ndim == 2
        logprob = torch.stack(
            [self.doublewell_logprob(x[:, i*2:i*2+2]) for i in range(self.n_wells)],
        dim=1).sum(dim=1)
        return logprob

    def sample_first_dimension(self, batch_size):
        def target_log_prob(x):
            return -x ** 4 + 6 * x ** 2 + 1 / 2 * x

        # Define proposal
        mix = torch.distributions.Categorical(self.component_mix)
        com = torch.distributions.Normal(self.means, self.scales)
        proposal = torch.distributions.MixtureSameFamily(mixture_distribution=mix,
                                                         component_distribution=com)

        k = self.Z_x1 * 3
        samples = rejection_sampling(batch_size, proposal, target_log_prob, k)
        return samples

    def sample_doublewell(self, batch_size):
        x1 = self.sample_first_dimension(batch_size)
        x2 = torch.randn_like(x1)
        return torch.stack([x1, x2], dim=1)

    def sample(self, batch_size):
        return torch.cat(
            [self.sample_doublewell(batch_size) for _ in range(self.n_wells)],
        dim=-1)

    def viz_pdf(self, samples=None, num_samples=5000):
        if samples is None:
            samples = self.sample(num_samples)

        x13 = samples[:, 0:3:2]
        viz_sample2d(x13, "samples", f"distx13.png", lim=3)
        viz_kde2d(x13, "kde", f"kdex13.png", lim=3)

        lim = 3
        alpha = 0.8
        n_contour_levels = 20
        def logp_func(x_2d):
            x = torch.zeros((x_2d.shape[0], self.data_ndim))
            x[:, 0] = x_2d[:, 0]
            x[:, 2] = x_2d[:, 1]
            return -self.energy(x)
        contour_img_path = f"contourx13.png"
        viz_contour_sample2d(x13, contour_img_path, logp_func, lim=lim, alpha=alpha, n_contour_levels=n_contour_levels)

        x23 = samples[:, 1:3]
        viz_sample2d(x23, "samples", f"distx23.png", lim=3)
        viz_kde2d(x23, "kde", f"kdex23.png", lim=3)

        def logp_func(x_2d):
            x = torch.zeros((x_2d.shape[0], self.data_ndim))
            x[:, 1] = x_2d[:, 0]
            x[:, 2] = x_2d[:, 1]
            return -self.energy(x)
        contour_img_path2 = f"contourx23.png"
        viz_contour_sample2d(x23, contour_img_path2, logp_func, lim=lim, alpha=alpha, n_contour_levels=n_contour_levels)

target/distribution/manywell.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
import torch as th


class RBFKernel:
    """
    A batch implementation
    """

    def __init__(self, beta=1e-2):
        self.beta = beta
        assert self.beta > 0

    def value(self, x, y):  # (B, 1)
        r = ((x - y) ** 2).sum(dim=-1, keepdim=True)
        return th.exp(-self.beta * r)

    def grad_x(self, x, y):  # (B, N)
        r = ((x - y) ** 2).sum(dim=-1, keepdim=True)
        return -2 * self.beta * (x - y) * th.exp(-self.beta * r)

    def grad_y(self, x, y):  # (B,N)
        r = ((x - y) ** 2).sum(dim=-1, keepdim=True)
        return 2 * self.beta * (x - y) * th.exp(-self.beta * r)

    def grad_xy(self, x, y):  # trace (B,N)
        assert len(x) == len(y)
        _, n = x.shape
        r = ((x - y) ** 2).sum(dim=-1, keepdim=True)  # (B,1)
        _y = 2 * self.beta * th.exp(-self.beta * r) * th.ones((1, n)).to(x)  # (B, N)
        _xy = 4 * self.beta ** 2 * (x - y) ** 2 * th.exp(-self.beta * r)
        return _y + _xy


class KSD:
    def __init__(self, score_fn, beta=1e-2):
        self.k_method = RBFKernel(beta)

        self.k = self.k_method.value
        self.grad_kx = self.k_method.grad_x
        self.grad_ky = self.k_method.grad_y
        self.grad_kxy = self.k_method.grad_xy

        self.score_fn = score_fn

    def cal_s(self, x, y):
        log_px = self.score_fn(x)
        log_py = self.score_fn(y)

        p1 = self.k(x, y) * (log_px * log_py).sum(
            dim=-1, keepdims=True
        )  # EQ: k(x,y) * S.T @ S, (B,1)
        p2 = (log_px * self.grad_ky(x, y)).sum(
            dim=-1, keepdims=True
        )  # EQ: Q.T @ \nabla_{x'}k, (B,1)
        p3 = (self.grad_kx(x, y) * log_py).sum(
            dim=-1, keepdims=True
        )  # EQ: \nabla_x k @ Q
        p4 = self.grad_kxy(x, y).sum(dim=-1, keepdims=True)

        return p1 + p2 + p3 + p4

    # calculate the KSD between score_fn and samples "x"
    def __call__(self, x, adjust_beta=False):
        # only support one batch dimension
        batch_size = x.shape[0]
        x = x.view(batch_size, -1)

        idx = th.combinations(th.arange(batch_size))
        x = x[idx]
        x, y = x[:, 0], x[:, 1]
        if adjust_beta:
            self.k_method.beta = 1.0 / (2 * (x - y) ** 2).mean()
        x, y = x.requires_grad_(True), y.requires_grad_(True)
        return (self.cal_s(x, y) + self.cal_s(y, x)).sum() / (
            batch_size * (batch_size - 1)
        )

target/ksd.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
import numpy as np
import torch


def nll_unit_gaussian(data, sigma=1.0):
    data = data.view(data.shape[0], -1)
    loss = 0.5 * np.log(2 * np.pi) + np.log(sigma) + 0.5 * data * data / (sigma ** 2)
    return torch.sum(torch.flatten(loss, start_dim=1), -1)

def linear_intepolate_energy(origin_energy_fn, x, weight_gauss=1.0):
    origin_energy = origin_energy_fn(x)
    gaussian_energy = nll_unit_gaussian(x)
    assert origin_energy.shape == gaussian_energy.shape
    return weight_gauss * gaussian_energy + (1.0 - weight_gauss) * origin_energy
target/loss_helper.py
zdhNarsil/Diffusion-Generative-Flow-Samplers
import itertools
import matplotlib.pyplot as plt
import seaborn as sns

import numpy as np
import torch
import torch as th
from einops import rearrange

# Inputs are all numpy arrays

########### 1D plot
def traj_plot1d(traj_len, samples, xlabel, ylabel, title="", fsave="img.png"):
    samples = rearrange(samples, "t b d -> b t d").cpu()
    inds = np.linspace(0, samples.shape[1], traj_len, endpoint=False, dtype=int)
    samples = samples[:, inds]
    plt.figure()
    for i, sample in enumerate(samples):
        plt.plot(np.arange(traj_len), sample.flatten(), marker="x", label=f"sample {i}")
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.savefig(fsave)
    plt.close()

def dist_plot1d(samples, nll_target_fn, nll_prior_fn=None,
                fname="img.png", width=8.):
    fig, ax = plt.subplots(1, 1, figsize=(7, 7))
    num_bins = 100
    density, bins = np.histogram(samples, num_bins, density=True)
    ax.plot(bins[1:], density, label="sampled")

    num_bins = 100
    query_x = th.linspace(-width/2, width/2, num_bins).cuda()
    target_unpdf = th.exp(-nll_target_fn(query_x.view(-1, 1)))
    target_norm_pdf = target_unpdf / th.sum(target_unpdf) / width * num_bins # make integral=1.
    np_x = query_x.cpu().numpy()
    ax.plot(np_x, target_norm_pdf.cpu().numpy(), label="target")

    if nll_prior_fn is not None:
        prior_unpdf = th.exp(-nll_prior_fn(query_x.view(-1, 1)))
        prior_norm_pdf = prior_unpdf / th.sum(prior_unpdf) / width * num_bins
        ax.plot(np_x, prior_norm_pdf.cpu().numpy(), label="prior")

    ax.set_xlim(np_x[0], np_x[-1])
    ax.set_ylim(0, 1.5 * th.max(target_norm_pdf).item())
    ax.legend()
    fig.savefig(fname)
    plt.close(fig)
    return fname

@torch.no_grad()
def drift_surface1d(model): # not used
    model.cuda()
    xs = th.linspace(-3.0, 3.0, 120).view(-1, 1).cuda()
    ts = th.linspace(0.0, 0.99, 100).cuda()

    values = []
    for cur_t in ts:
        values.append(model.f_func(cur_t, xs))
    values = th.cat(values, dim=1)

    x, t, zz = map(lambda x: x.cpu().numpy(), [xs, ts, values])
    tt, xx = np.meshgrid(t, x)
    return tt, xx, zz


########### 2D plot
def viz_sample2d(points, title, fsave, lim=7.0, sample_num=50000):
    fig, axs = plt.subplots(1, 1, figsize=(7, 7))
    if title is not None:
        axs.set_title(title)
    axs.plot(
        points[:sample_num, 0],
        points[:sample_num, 1],
        linewidth=0,
        marker=".",
        markersize=1,
    )
    axs.set_xlim(-lim, lim)
    axs.set_ylim(-lim, lim)
    if "pdf" in fsave:
        fig.savefig(fsave, format="pdf", bbox_inches='tight')
    else:
        fig.savefig(fsave, bbox_inches='tight')
    plt.close(fig)
    plt.close()
    plt.cla()
    plt.clf()

def viz_kde2d(points, title, fname, lim=7.0, sample_num=2000):
    fig, ax = plt.subplots(1, 1, figsize=(7, 7), dpi=200)
    if title is not None:
        ax.set_title(title)
    sns.kdeplot(
        x=points[:sample_num, 0], y=points[:sample_num, 1],
        cmap="coolwarm", fill=True, ax=ax
    )
    ax.set_xlim(-lim, lim)
    ax.set_ylim(-lim, lim)
    # ax.axis("off")
    if "pdf" in fname:
        fig.savefig(fname, format="pdf", bbox_inches='tight')
    else:
        fig.savefig(fname, bbox_inches='tight')
    plt.close(fig)


def viz_coutour_with_ax(ax, log_prob_func, lim=3.0, n_contour_levels=None):
    grid_width_n_points = 100
    log_prob_min = -1000.0
    # plot_contours
    x_points_dim1 = torch.linspace(-lim, lim, grid_width_n_points)
    x_points_dim2 = x_points_dim1
    x_points = torch.tensor(list(itertools.product(x_points_dim1, x_points_dim2)))
    log_p_x = log_prob_func(x_points).detach()
    log_p_x = torch.clamp_min(log_p_x, log_prob_min)
    log_p_x = log_p_x.reshape((grid_width_n_points, grid_width_n_points))
    x_points_dim1 = x_points[:, 0].reshape((grid_width_n_points, grid_width_n_points)).numpy()
    x_points_dim2 = x_points[:, 1].reshape((grid_width_n_points, grid_width_n_points)).numpy()
    if n_contour_levels:
        ax.contour(x_points_dim1, x_points_dim2, log_p_x, levels=n_contour_levels)
    else:
        ax.contour(x_points_dim1, x_points_dim2, log_p_x)


def viz_contour_sample2d(points, fname, log_prob_func,
                         lim=3.0, alpha=0.7, n_contour_levels=None):
    # plotting_bounds = (-3, 3) # for manywells
    fig, ax = plt.subplots(1, 1, figsize=(7, 7))

    viz_coutour_with_ax(ax, log_prob_func, lim=lim, n_contour_levels=n_contour_levels)

    # plot samples
    samples = torch.clamp(points, -lim, lim)
    samples = samples.cpu().detach()
    ax.plot(samples[:, 0], samples[:, 1],
            linewidth=0, marker=".", markersize=1.5, alpha=alpha)

    # ax.set_xlim(-lim, lim)
    # ax.set_ylim(-lim, lim)
    fig.savefig(fname, bbox_inches='tight')
    plt.close(fig)
target/plot.py
zdhNarsil/Diffusion-Generative-Flow-Samplers