Add files using upload-large-folder tool

f8faa82 verified 18 days ago

8.62 kB

from itertools import cycle
import math
import os
import torch
import tyro
import numpy as np
from dataclasses import dataclass, field

from datasets.sequence import SequenceDataset
from model import TemporalUnet, CosineScheduleDiffusion, apply_conditioning
from utils.arrays import batch_to_device, to_np, to_torch, apply_dict
from utils.training import deterministic, set_cuda_visible_device, save_config
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
import tqdm
import wandb

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as patches


@dataclass
class TrainingConfig:
    seed: int = 0
    device: str = "cuda:0"
    log_folder: str = "logs"
    exp_name: str = "diffusion"

    env: str = "avoiding-v0"
    horizon: int = 16
    normalizer: str = "LimitsNormalizer"
    preprocess_fns: list = field(default_factory=lambda: [])
    max_path_length: int = 200
    max_n_episodes: int = 500
    termination_penalty: float = 0

    state_dim: int = 4
    action_dim: int = 2

    # training
    n_train_steps: int = 1000001
    save_freq: int = 50000
    batch_size: int = 32
    learning_rate: float = 2e-4

    ema_decay: float = 0.995

    sample_freq: int = 5000
    n_sample_steps: int = 20
    n_samples: int = 8

    resume: bool = False

    wandb_project: str = "direct-d3il"
    wandb_entity: str = ""


@torch.no_grad()
def generate_samples(ema_model, diffusion, normalizer, dataset, cfg):
    """Generate DDIM samples from the EMA model for visualization (no conditioning)."""
    ema_model.eval()
    action_dim = cfg.action_dim
    state_dim = cfg.state_dim
    horizon = cfg.horizon
    transition_dim = action_dim + state_dim
    device = cfg.device

    # grab a real initial observation from the dataset
    sample = dataset[0]
    obs_0 = sample.conditions[0]  # already normalized
    conditions = {0: obs_0}

    x = torch.randn(cfg.n_samples, horizon, transition_dim, device=device)
    x = apply_conditioning(x, to_torch(conditions, device=device), action_dim)

    timesteps = torch.linspace(1.0, 0.0, cfg.n_sample_steps + 1, device=device)

    for i in range(cfg.n_sample_steps):
        t_now = timesteps[i]
        t_next = timesteps[i + 1]

        abar_now = diffusion.alpha_bar(t_now)
        abar_next = diffusion.alpha_bar(t_next)

        t_batch = t_now.expand(cfg.n_samples)
        x0_pred = ema_model(x, t_batch)

        # DDIM (deterministic): x_{t_next} = √ᾱ_{t_next} * x0_pred + √(1-ᾱ_{t_next}) * ε_pred
        # where ε_pred = (x - √ᾱ_t * x0_pred) / √(1-ᾱ_t)
        eps_pred = (x - math.sqrt(abar_now) * x0_pred) / math.sqrt(1.0 - abar_now)
        x = math.sqrt(abar_next) * x0_pred + math.sqrt(1.0 - abar_next) * eps_pred

        # only condition the start point (not the endpoint)
        x = apply_conditioning(x, to_torch(conditions, device=device), action_dim)

    # unnormalize observations for plotting
    obs = normalizer.unnormalize(to_np(x[:, :, action_dim:]), "observations")
    return obs  # (n_samples, horizon, state_dim)


def plot_samples(observations, save_path):
    """Plot sampled trajectories on the obstacle field."""
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_xlim([0.2, 0.8])
    ax.set_ylim([-0.3, 0.4])
    ax.set_facecolor([1, 1, 0.9])
    ax.set_aspect("equal", adjustable="box")
    ax.grid(True, alpha=0.3)

    # target
    ax.add_patch(patches.Rectangle((0.2, 0.35), 0.6, 0.05,
                                    facecolor=[0.4, 1, 0.4], alpha=0.35))
    # pillars
    for c, r in [([0.5,-0.1],0.03), ([0.425,0.08],0.025), ([0.575,0.08],0.025),
                  ([0.35,0.26],0.025), ([0.5,0.26],0.025), ([0.65,0.26],0.025)]:
        ax.add_patch(patches.Circle(c, r, color="r", alpha=0.7))

    colors = plt.cm.tab10(np.linspace(0, 1, len(observations)))
    for i, obs in enumerate(observations):
        x_pos = obs[:, 2]  # actual_x
        y_pos = obs[:, 3]  # actual_y
        ax.plot(x_pos, y_pos, "-o", color=colors[i], linewidth=1.5, markersize=3, alpha=0.8)
        ax.plot(x_pos[0], y_pos[0], "s", color="green", markersize=5)

    plt.tight_layout()
    fig.savefig(save_path, dpi=100, bbox_inches="tight")
    plt.close(fig)


def train(cfg: TrainingConfig, log_subfolder: str):
    dataset = SequenceDataset(
        env=cfg.env,
        horizon=cfg.horizon,
        normalizer=cfg.normalizer,
        preprocess_fns=cfg.preprocess_fns,
        max_path_length=cfg.max_path_length,
        max_n_episodes=cfg.max_n_episodes,
        termination_penalty=0,
        seed=cfg.seed,
    )

    unet = TemporalUnet(
        horizon=cfg.horizon,
        transition_dim=cfg.state_dim + cfg.action_dim,
        cond_dim=cfg.state_dim,
        dim=32,
        dim_mults=(1, 4, 8),
    ).to(cfg.device)

    diffusion = CosineScheduleDiffusion(
        model=unet,
        action_dim=cfg.action_dim,
    )

    train_loader = cycle(DataLoader(dataset, batch_size=cfg.batch_size, shuffle=True))

    optimizer = torch.optim.Adam(unet.parameters(), lr=cfg.learning_rate)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=cfg.save_freq * 2
    )
    ema = torch.optim.swa_utils.AveragedModel(
        unet,
        avg_fn=lambda avg, new, num: cfg.ema_decay * avg + (1 - cfg.ema_decay) * new,
    )

    # resume from latest checkpoint
    start_step = 0
    if cfg.resume:
        import glob
        ckpts = sorted(glob.glob(os.path.join(log_subfolder, "model_*.pth")))
        ckpts = [c for c in ckpts if "ema" not in c]
        if ckpts:
            latest = ckpts[-1]
            cp_idx = int(latest.split("model_")[-1].replace(".pth", ""))
            start_step = cp_idx * cfg.save_freq

            unet.load_state_dict(torch.load(latest))
            ema_path = latest.replace("model_", "model_ema_")
            if os.path.exists(ema_path):
                ema.module.load_state_dict(torch.load(ema_path))

            opt_path = os.path.join(log_subfolder, "optimizer.pth")
            if os.path.exists(opt_path):
                opt_state = torch.load(opt_path)
                optimizer.load_state_dict(opt_state["optimizer"])
                scheduler.load_state_dict(opt_state["scheduler"])

            print(f"Resumed from checkpoint {cp_idx} (step {start_step})")
        else:
            print("No checkpoints found, starting from scratch")

    writer = SummaryWriter(log_dir=os.path.join(log_subfolder, "tensorboard_logs"))

    wandb.init(
        project=cfg.wandb_project,
        entity=cfg.wandb_entity or None,
        name=cfg.exp_name,
        config=vars(cfg),
        resume="allow" if cfg.resume else None,
    )

    samples_dir = os.path.join(log_subfolder, "samples")
    os.makedirs(samples_dir, exist_ok=True)

    # skip dataloader to the right position
    if start_step > 0:
        print(f"Skipping dataloader to step {start_step}...")
        for _ in tqdm.tqdm(range(start_step), desc="Fast-forwarding"):
            next(train_loader)

    for i in tqdm.tqdm(range(start_step, cfg.n_train_steps)):
        batch = batch_to_device(next(train_loader), cfg.device)
        loss, infos = diffusion.loss(*batch)
        loss.backward()
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

        ema.update_parameters(unet)

        if i % cfg.save_freq == 0:
            torch.save(
                unet.state_dict(),
                os.path.join(log_subfolder, f"model_{i // cfg.save_freq}.pth"),
            )
            torch.save(
                ema.module.state_dict(),
                os.path.join(log_subfolder, f"model_ema_{i // cfg.save_freq}.pth"),
            )
            torch.save(
                {"optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict()},
                os.path.join(log_subfolder, "optimizer.pth"),
            )

        if i % cfg.sample_freq == 0:
            obs = generate_samples(ema.module, diffusion, dataset.normalizer, dataset, cfg)
            sample_path = os.path.join(samples_dir, f"step_{i:07d}.png")
            plot_samples(obs, sample_path)
            wandb.log({"samples": wandb.Image(sample_path)}, step=i)

        writer.add_scalar("loss", loss, i)
        wandb.log({"loss": loss.item(), "lr": scheduler.get_last_lr()[0]}, step=i)

    wandb.finish()


if __name__ == "__main__":
    cfg = tyro.cli(TrainingConfig)

    set_cuda_visible_device(cfg)
    deterministic(cfg.seed)

    log_subfolder = os.path.join(cfg.log_folder, cfg.env, "diffusion", cfg.exp_name)
    save_config(cfg, log_subfolder)

    train(cfg, log_subfolder)

	from itertools import cycle
	import math
	import os
	import torch
	import tyro
	import numpy as np
	from dataclasses import dataclass, field

	from datasets.sequence import SequenceDataset
	from model import TemporalUnet, CosineScheduleDiffusion, apply_conditioning
	from utils.arrays import batch_to_device, to_np, to_torch, apply_dict
	from utils.training import deterministic, set_cuda_visible_device, save_config
	from torch.utils.tensorboard import SummaryWriter
	from torch.utils.data import DataLoader
	import tqdm
	import wandb

	import matplotlib
	matplotlib.use("Agg")
	import matplotlib.pyplot as plt
	import matplotlib.patches as patches


	@dataclass
	class TrainingConfig:
	seed: int = 0
	device: str = "cuda:0"
	log_folder: str = "logs"
	exp_name: str = "diffusion"

	env: str = "avoiding-v0"
	horizon: int = 16
	normalizer: str = "LimitsNormalizer"
	preprocess_fns: list = field(default_factory=lambda: [])
	max_path_length: int = 200
	max_n_episodes: int = 500
	termination_penalty: float = 0

	state_dim: int = 4
	action_dim: int = 2

	# training
	n_train_steps: int = 1000001
	save_freq: int = 50000
	batch_size: int = 32
	learning_rate: float = 2e-4

	ema_decay: float = 0.995

	sample_freq: int = 5000
	n_sample_steps: int = 20
	n_samples: int = 8

	resume: bool = False

	wandb_project: str = "direct-d3il"
	wandb_entity: str = ""


	@torch.no_grad()
	def generate_samples(ema_model, diffusion, normalizer, dataset, cfg):
	"""Generate DDIM samples from the EMA model for visualization (no conditioning)."""
	ema_model.eval()
	action_dim = cfg.action_dim
	state_dim = cfg.state_dim
	horizon = cfg.horizon
	transition_dim = action_dim + state_dim
	device = cfg.device

	# grab a real initial observation from the dataset
	sample = dataset[0]
	obs_0 = sample.conditions[0] # already normalized
	conditions = {0: obs_0}

	x = torch.randn(cfg.n_samples, horizon, transition_dim, device=device)
	x = apply_conditioning(x, to_torch(conditions, device=device), action_dim)

	timesteps = torch.linspace(1.0, 0.0, cfg.n_sample_steps + 1, device=device)

	for i in range(cfg.n_sample_steps):
	t_now = timesteps[i]
	t_next = timesteps[i + 1]

	abar_now = diffusion.alpha_bar(t_now)
	abar_next = diffusion.alpha_bar(t_next)

	t_batch = t_now.expand(cfg.n_samples)
	x0_pred = ema_model(x, t_batch)

	# DDIM (deterministic): x_{t_next} = √ᾱ_{t_next} * x0_pred + √(1-ᾱ_{t_next}) * ε_pred
	# where ε_pred = (x - √ᾱ_t * x0_pred) / √(1-ᾱ_t)
	eps_pred = (x - math.sqrt(abar_now) * x0_pred) / math.sqrt(1.0 - abar_now)
	x = math.sqrt(abar_next) * x0_pred + math.sqrt(1.0 - abar_next) * eps_pred

	# only condition the start point (not the endpoint)
	x = apply_conditioning(x, to_torch(conditions, device=device), action_dim)

	# unnormalize observations for plotting
	obs = normalizer.unnormalize(to_np(x[:, :, action_dim:]), "observations")
	return obs # (n_samples, horizon, state_dim)


	def plot_samples(observations, save_path):
	"""Plot sampled trajectories on the obstacle field."""
	fig, ax = plt.subplots(figsize=(6, 6))
	ax.set_xlim([0.2, 0.8])
	ax.set_ylim([-0.3, 0.4])
	ax.set_facecolor([1, 1, 0.9])
	ax.set_aspect("equal", adjustable="box")
	ax.grid(True, alpha=0.3)

	# target
	ax.add_patch(patches.Rectangle((0.2, 0.35), 0.6, 0.05,
	facecolor=[0.4, 1, 0.4], alpha=0.35))
	# pillars
	for c, r in [([0.5,-0.1],0.03), ([0.425,0.08],0.025), ([0.575,0.08],0.025),
	([0.35,0.26],0.025), ([0.5,0.26],0.025), ([0.65,0.26],0.025)]:
	ax.add_patch(patches.Circle(c, r, color="r", alpha=0.7))

	colors = plt.cm.tab10(np.linspace(0, 1, len(observations)))
	for i, obs in enumerate(observations):
	x_pos = obs[:, 2] # actual_x
	y_pos = obs[:, 3] # actual_y
	ax.plot(x_pos, y_pos, "-o", color=colors[i], linewidth=1.5, markersize=3, alpha=0.8)
	ax.plot(x_pos[0], y_pos[0], "s", color="green", markersize=5)

	plt.tight_layout()
	fig.savefig(save_path, dpi=100, bbox_inches="tight")
	plt.close(fig)


	def train(cfg: TrainingConfig, log_subfolder: str):
	dataset = SequenceDataset(
	env=cfg.env,
	horizon=cfg.horizon,
	normalizer=cfg.normalizer,
	preprocess_fns=cfg.preprocess_fns,
	max_path_length=cfg.max_path_length,
	max_n_episodes=cfg.max_n_episodes,
	termination_penalty=0,
	seed=cfg.seed,
	)

	unet = TemporalUnet(
	horizon=cfg.horizon,
	transition_dim=cfg.state_dim + cfg.action_dim,
	cond_dim=cfg.state_dim,
	dim=32,
	dim_mults=(1, 4, 8),
	).to(cfg.device)

	diffusion = CosineScheduleDiffusion(
	model=unet,
	action_dim=cfg.action_dim,
	)

	train_loader = cycle(DataLoader(dataset, batch_size=cfg.batch_size, shuffle=True))

	optimizer = torch.optim.Adam(unet.parameters(), lr=cfg.learning_rate)
	scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
	optimizer, T_max=cfg.save_freq * 2
	)
	ema = torch.optim.swa_utils.AveragedModel(
	unet,
	avg_fn=lambda avg, new, num: cfg.ema_decay * avg + (1 - cfg.ema_decay) * new,
	)

	# resume from latest checkpoint
	start_step = 0
	if cfg.resume:
	import glob
	ckpts = sorted(glob.glob(os.path.join(log_subfolder, "model_*.pth")))
	ckpts = [c for c in ckpts if "ema" not in c]
	if ckpts:
	latest = ckpts[-1]
	cp_idx = int(latest.split("model_")[-1].replace(".pth", ""))
	start_step = cp_idx * cfg.save_freq

	unet.load_state_dict(torch.load(latest))
	ema_path = latest.replace("model_", "model_ema_")
	if os.path.exists(ema_path):
	ema.module.load_state_dict(torch.load(ema_path))

	opt_path = os.path.join(log_subfolder, "optimizer.pth")
	if os.path.exists(opt_path):
	opt_state = torch.load(opt_path)
	optimizer.load_state_dict(opt_state["optimizer"])
	scheduler.load_state_dict(opt_state["scheduler"])

	print(f"Resumed from checkpoint {cp_idx} (step {start_step})")
	else:
	print("No checkpoints found, starting from scratch")

	writer = SummaryWriter(log_dir=os.path.join(log_subfolder, "tensorboard_logs"))

	wandb.init(
	project=cfg.wandb_project,
	entity=cfg.wandb_entity or None,
	name=cfg.exp_name,
	config=vars(cfg),
	resume="allow" if cfg.resume else None,
	)

	samples_dir = os.path.join(log_subfolder, "samples")
	os.makedirs(samples_dir, exist_ok=True)

	# skip dataloader to the right position
	if start_step > 0:
	print(f"Skipping dataloader to step {start_step}...")
	for _ in tqdm.tqdm(range(start_step), desc="Fast-forwarding"):
	next(train_loader)

	for i in tqdm.tqdm(range(start_step, cfg.n_train_steps)):
	batch = batch_to_device(next(train_loader), cfg.device)
	loss, infos = diffusion.loss(*batch)
	loss.backward()
	optimizer.step()
	scheduler.step()
	optimizer.zero_grad()

	ema.update_parameters(unet)

	if i % cfg.save_freq == 0:
	torch.save(
	unet.state_dict(),
	os.path.join(log_subfolder, f"model_{i // cfg.save_freq}.pth"),
	)
	torch.save(
	ema.module.state_dict(),
	os.path.join(log_subfolder, f"model_ema_{i // cfg.save_freq}.pth"),
	)
	torch.save(
	{"optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict()},
	os.path.join(log_subfolder, "optimizer.pth"),
	)

	if i % cfg.sample_freq == 0:
	obs = generate_samples(ema.module, diffusion, dataset.normalizer, dataset, cfg)
	sample_path = os.path.join(samples_dir, f"step_{i:07d}.png")
	plot_samples(obs, sample_path)
	wandb.log({"samples": wandb.Image(sample_path)}, step=i)

	writer.add_scalar("loss", loss, i)
	wandb.log({"loss": loss.item(), "lr": scheduler.get_last_lr()[0]}, step=i)

	wandb.finish()


	if __name__ == "__main__":
	cfg = tyro.cli(TrainingConfig)

	set_cuda_visible_device(cfg)
	deterministic(cfg.seed)

	log_subfolder = os.path.join(cfg.log_folder, cfg.env, "diffusion", cfg.exp_name)
	save_config(cfg, log_subfolder)

	train(cfg, log_subfolder)