import torch
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from diffusion_policy.model.common.lr_scheduler import get_scheduler


epochs = 2000
len_dataset = 10000

params = torch.Tensor(1, 1, 1, 1)
optimizer = AdamW([params], lr=1.0e-4, betas=[0.95, 0.999], eps=1.0e-8, weight_decay=1.0e-6)
lr_scheduler: LambdaLR = get_scheduler(
    "cosine",
    optimizer=optimizer,
    num_warmup_steps=500,
    num_training_steps=(len_dataset * epochs),
    # pytorch assumes stepping LRScheduler every epoch
    # however huggingface diffusers steps it every batch
)

for epoch in range(epochs):
    # for _ in range(len_dataset):
    lr_scheduler.step()
    # print(lr_scheduler.get_last_lr())

    if epoch % 100 == 0:
        print(f"Epoch: {epoch}, LR: {lr_scheduler.get_last_lr()}")