Adding score function code (and U-Net code) for pixel single-integrator.

examples/pixels_singleint/config/train.yaml

@@ -0,0 +1,37 @@
+defaults:
+  - user: glen
+
+data:
+  channels: 2
+
+hydra:
+  run:
+    dir: ${user.run_dir}
+
+nn_layers: [512, 512, 512, 512]
+dataset_size: 100000
+
+load_data: True
+load_data_path: '/home/gchou/PycharmProjects/score_po_data/pixels_singleint/100000.pt'
+
+model:
+  ngf: 128
+  num_classes: 232
+  sigma_begin: 50
+  sigma_end: 0.01
+
+train:
+  wandb:
+    enabled: ${user.wandb_enabled}
+    project: score_po
+    dir: ${user.run_dir}
+    name: ${now:%Y.%m.%d-%H.%M.%S}_pixels_singleint
+    entity: ${user.wandb_entity}
+  adam:
+    lr: 1e-4
+    epochs: 2000
+    batch_size: 128
+  dataset_split: [0.99, 0.01]
+  save_best_model: ncsn_xux.pth
+  load_ckpt: null
+  device: "cuda"

examples/pixels_singleint/config/user/glen.yaml

@@ -0,0 +1,4 @@
+run_dir: /home/gchou/PycharmProjects/score_po_data/${now:%Y-%m-%d}/${now:%H-%M-%S}
+
+wandb_enabled: True
+wandb_entity: gchou

examples/pixels_singleint/dynamics.py

@@ -0,0 +1,16 @@
+from score_po.dynamical_system import DynamicalSystem
+
+
+class SingleIntegratorPixels(DynamicalSystem):
+    def __init__(self):
+        super().__init__(2, 2)
+        self.is_differentiable = True
+
+    def dynamics(self, x, u):
+        return x + u
+
+    def dynamics_batch(self, x_batch, u_batch):
+        return x_batch + u_batch
+
+
+dynamics = SingleIntegratorPixels()

examples/pixels_singleint/environment.py

@@ -0,0 +1,77 @@
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from PIL import Image, ImageDraw, ImageFilter
+
+class Environment:
+    """
+    Generate a 2D environment with obstacles.
+    """
+
+    def __init__(self):
+        self.shape_lst = []
+
+    def add_ellipse(self, mu, a, b):
+        self.shape_lst.append(["ellipse", mu, a, b])
+
+    def add_rectangle(self, mu, a, b):
+        self.shape_lst.append(["rectangle", mu, a, b])
+
+    def sample_ellipse(self, mu, a, b, samples):
+        return ((samples[:, 0] - mu[0]) / a) ** 2 + ((samples[:, 1] - mu[1]) / b) ** 2 < 1.0
+
+    def sample_rectangle(self, mu, a, b, samples):
+        return torch.logical_and(
+            (torch.abs(samples[:, 0] - mu[0]) < a / 2),
+            (torch.abs(samples[:, 1] - mu[1]) < b / 2)
+        )
+
+    def sample_image(self, num_points, dim, samples=None, filter_obs=True, xlim=1., ell_size=0.05, N_EXP=10, buf=0.2):
+        if samples is None:
+            samples = 2.0 * xlim * torch.rand(num_points, 2) - xlim
+            valid = torch.zeros(num_points)
+        else:
+            valid = torch.zeros(samples.shape[0])
+        robot_size = ell_size*dim*N_EXP
+
+        samples_img = []
+        for i in range(samples.shape[0]):
+            x = samples[i, 0]
+            y = samples[i, 1]
+            shape = [(x + (1+buf)*xlim)/(2*(1+buf)*xlim)*dim*N_EXP - robot_size,
+                     (y + (1+buf)*xlim)/(2*(1+buf)*xlim)*dim*N_EXP - robot_size,
+                     (x + (1+buf)*xlim)/(2*(1+buf)*xlim)*dim*N_EXP + robot_size,
+                     (y + (1+buf)*xlim)/(2*(1+buf)*xlim)*dim*N_EXP + robot_size]
+
+            # creating new Image object
+            img = Image.new("RGB", (dim*N_EXP, dim*N_EXP))
+            # create circle robot
+            img1 = ImageDraw.Draw(img)
+            img1.ellipse(shape, fill="red", outline="red")
+            img = np.array(img)
+            img = Image.fromarray(img).filter(ImageFilter.BoxBlur(N_EXP))
+            img = np.array(img.resize((int(dim), int(dim)), resample=Image.BICUBIC))
+
+            samples_img.append(img[:, :, 0] / img[:, :, 0].max())
+        if filter_obs == True:
+            for shape in self.shape_lst:
+                if shape[0] == "ellipse":
+                    valid = torch.logical_or(valid,
+                                             self.sample_ellipse(shape[1], shape[2], shape[3], samples)
+                                             )
+                if shape[0] == "rectangle":
+                    valid = torch.logical_or(valid,
+                                             self.sample_rectangle(shape[1], shape[2], shape[3], samples)
+                                             )
+        samples_img = torch.tensor(samples_img)
+        return samples_img[~valid.bool(), :], samples[~valid.bool(), :]
+
+
+def test():
+    env = Environment()
+    env.add_ellipse([0, 0], 0.4, 0.4)
+    pts = env.sample_image(100000)
+    plt.figure()
+    plt.plot(pts[:, 0], pts[:, 1], 'ro')
+    plt.savefig("env.png")
+    plt.close()

examples/pixels_singleint/scripts/train_ncsn_xux.py

@@ -0,0 +1,51 @@
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from torch.utils.data import TensorDataset
+
+import hydra
+from omegaconf import DictConfig
+
+from score_po.score_matching import NoiseConditionedScoreEstimatorXux, NoiseConditionedScoreEstimatorXuxImage
+from score_po.nn import (
+    MLPwEmbedding, NCSNv2, TrainParams, Normalizer, generate_cosine_schedule, get_sigmas)
+
+from examples.pixels_singleint.dynamics import SingleIntegratorPixels
+from examples.pixels_singleint.environment import Environment
+
+
+@hydra.main(config_path="../config", config_name="train")
+def main(cfg: DictConfig):
+    env = Environment()
+    dynamics = SingleIntegratorPixels()
+    env.add_ellipse([0, 0], 0.4, 0.4)
+    nx = 32  # image width/height
+    nu = 2  # number of control inputs
+    u_normalizer = Normalizer(k=0.2 * torch.ones(2), b=torch.zeros(2))
+    if not cfg.load_data:
+        y_batch, x_batch = env.sample_image(cfg.dataset_size, nx, filter_obs=True)
+        u_batch = 0.1 * 2.0 * (torch.rand(x_batch.shape[0], 2) - 0.5)
+        xnext_batch = dynamics.dynamics_batch(x_batch, u_batch)
+        ynext_batch, _ = env.sample_image(cfg.dataset_size, nx, samples=xnext_batch, filter_obs=False)
+    else:
+        y_batch, u_batch, ynext_batch = torch.load(cfg.load_data_path)
+        # y_batch = y_batch[:10000]
+        # u_batch = u_batch[:10000]
+        # ynext_batch = ynext_batch[:10000]
+    dataset = TensorDataset(y_batch.reshape(y_batch.shape[0], -1).float().to(cfg.train.device),
+                            u_batch.float().to(cfg.train.device),
+                            ynext_batch.reshape(y_batch.shape[0], -1).float().to(cfg.train.device)
+                            )
+
+    params = TrainParams()
+    params.load_from_config(cfg)
+
+    # network = MLPwEmbedding(2*int(nx**2) + nu, 2*int(nx**2) + nu, cfg.nn_layers, 10)
+    network = NCSNv2(cfg, nu, nu, cfg.nn_layers).to(cfg.train.device)
+    sf = NoiseConditionedScoreEstimatorXuxImage(
+        int(nx**2), nu, get_sigmas(cfg), network, x_normalizer=None, u_normalizer=u_normalizer)
+    sf.train_network(dataset, params, sigma_lst=get_sigmas(cfg))
+
+
+if __name__ == "__main__":
+    main()

examples/pixels_singleint/scripts/train_score_xux.py

@@ -0,0 +1,47 @@
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from torch.utils.data import TensorDataset
+
+import hydra
+from omegaconf import DictConfig
+
+from score_po.score_matching import ScoreEstimatorXux
+from score_po.nn import MLP, TrainParams, Normalizer
+
+from examples.pixels_singleint.dynamics import SingleIntegratorPixels
+from examples.pixels_singleint.environment import Environment
+
+
+@hydra.main(config_path="../config", config_name="train")
+def main(cfg: DictConfig):
+    env = Environment()
+    dynamics = SingleIntegratorPixels()
+    env.add_ellipse([0, 0], 0.4, 0.4)
+    nx = 32 # image dimension
+    nu = 2 # control dimension
+
+    u_normalizer = Normalizer(k=0.2 * torch.ones(2), b=torch.zeros(2))
+    if not cfg.load_data:
+        y_batch, x_batch = env.sample_image(cfg.dataset_size, nx, filter_obs=True)
+        u_batch = 0.1 * 2.0 * (torch.rand(x_batch.shape[0], 2) - 0.5)
+        xnext_batch = dynamics.dynamics_batch(x_batch, u_batch)
+        ynext_batch, _ = env.sample_image(cfg.dataset_size, nx, samples=xnext_batch, filter_obs=False)
+    else:
+        y_batch, u_batch, ynext_batch = torch.load(cfg.load_data_path)
+    dataset = TensorDataset(y_batch.reshape(y_batch.shape[0], -1).float().to(cfg.device),
+                            u_batch.float().to(cfg.device),
+                            ynext_batch.reshape(y_batch.shape[0], -1).float().to(cfg.device)
+                            )
+
+    params = TrainParams()
+    params.load_from_config(cfg)
+
+    network = MLP(2*int(nx**2) + nu, 2*int(nx**2) + nu, cfg.nn_layers)
+    sf = ScoreEstimatorXux(
+        int(nx**2), nu, network, x_normalizer=None, u_normalizer=u_normalizer)
+    sf.train_network(dataset, params, sigma=0.1 * torch.ones(1))
+
+
+if __name__ == "__main__":
+    main()

examples/pixels_singleint/scripts/trajopt_diffusion_ncse.py

@@ -0,0 +1,85 @@
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from matplotlib.patches import Circle
+from matplotlib.pyplot import get_cmap
+import os
+
+import hydra
+from hydra.utils import get_original_cwd
+from omegaconf import DictConfig
+
+from score_po.score_matching import NoiseConditionedScoreEstimatorXuxImage
+from score_po.nn import MLPwEmbedding, TrainParams, NCSNv2, get_sigmas
+from score_po.trajectory_optimizer import (
+    TrajectoryOptimizer,
+    TrajectoryOptimizerNCSF, TrajectoryOptimizerSFParams)
+from score_po.trajectory import BVPTrajectory
+from score_po.costs import QuadraticCostImage
+
+from examples.light_dark.dynamics import SingleIntegrator
+from examples.light_dark.environment import Environment
+
+
+@hydra.main(config_path="../config", config_name="trajopt_diffusion")
+def main(cfg: DictConfig):
+    # 1. Set up parameters.
+    params = TrajectoryOptimizerSFParams()
+    nx = 32
+    nu = 2
+
+    # 2. Load score function.
+    network = NCSNv2(cfg, nu, nu, cfg.nn_layers)
+    sf = NoiseConditionedScoreEstimatorXuxImage(int(nx**2), nu, get_sigmas(cfg), network)
+    sf.load_state_dict(torch.load(os.path.join(
+        get_original_cwd(), "examples/pixels_singleint/weights/ncsn_xux.pth"
+    )))
+    params.sf = sf
+
+    # 3. Load costs.
+    cost = QuadraticCostImage()
+    cost.load_from_config(cfg)
+    params.cost = cost
+
+    # 4. Set up trajectory.
+    trj = BVPTrajectory(int(nx**2), nu, cfg.trj.T,
+                        torch.Tensor(cfg.trj.x0), torch.Tensor(cfg.trj.xT))
+    params.trj = trj
+
+    # 4. Set up optimizer
+    params.load_from_config(cfg)
+    params.to_device(cfg.trj.device)
+
+    # 5. Define callback function
+
+    def callback(params: TrajectoryOptimizerSFParams, loss: float, iter: int):
+        if iter % cfg.plot_period == 0:
+            x_trj, u_trj = params.trj.get_full_trajectory()
+            x_trj = x_trj.detach().cpu().numpy()
+            u_trj = u_trj.detach().cpu().numpy()
+
+            colormap = get_cmap('winter')
+
+            plt.figure(figsize=(8, 8))
+            for t in range(cfg.trj.T + 1):
+                plt.plot(x_trj[t, 0], x_trj[t, 1], marker='o', color=colormap(
+                    t / (cfg.trj.T + 1)
+                ))
+            for t in range(cfg.trj.T):
+                plt.arrow(x_trj[t, 0], x_trj[t, 1], u_trj[t, 0], u_trj[t, 1], color=colormap(
+                    t / cfg.trj.T))
+            circle = Circle([0, 0], 0.4, fill=True, color='k')
+            plt.gca().add_patch(circle)
+            plt.axis('equal')
+            plt.xlim([-1, 1])
+            plt.ylim([-1, 1])
+            plt.savefig("{:04d}.png".format(iter))
+            plt.close()
+
+    # 5. Run
+    optimizer = TrajectoryOptimizerNCSF(params)
+    optimizer.iterate(callback)
+
+
+if __name__ == "__main__":
+    main()

examples/pixels_singleint_line/config/train.yaml

@@ -0,0 +1,37 @@
+defaults:
+  - user: user_a
+
+data:
+  channels: 2
+
+hydra:
+  run:
+    dir: ${user.run_dir}
+
+nn_layers: [512, 512, 512, 512]
+dataset_size: 200000
+
+load_data: True
+load_data_path: '/home/Documents/PycharmProjects/score_po_data/pixels_singleint_line/200000.pt'
+
+model:
+  ngf: 128
+  num_classes: 232
+  sigma_begin: 50
+  sigma_end: 0.01
+
+train:
+  wandb:
+    enabled: ${user.wandb_enabled}
+    project: score_po
+    dir: ${user.run_dir}
+    name: ${now:%Y.%m.%d-%H.%M.%S}_pixels_singleint_line
+    entity: ${user.wandb_entity}
+  adam:
+    lr: 1e-4
+    epochs: 2000
+    batch_size: 128
+  dataset_split: [0.99, 0.01]
+  save_best_model: ncsn_xux.pth
+  load_ckpt: null
+  device: "cuda"