sweep.py

"""
Usage:
Training:
python train.py --config-name=example
"""
import logging
import os
import pathlib

import hydra
import wandb
import torch

from eval import eval_main
from omegaconf import DictConfig, OmegaConf

log = logging.getLogger(__name__)

OmegaConf.register_new_resolver("eval", eval, replace=True)


@hydra.main(
        version_base=None,
        config_path=str(pathlib.Path(__file__).parent.joinpath('imitation','config')),
        config_name="train" # use train parameters as base
        )
def train(cfg: DictConfig) -> None:
    print(OmegaConf.to_yaml(cfg))
    log.info("Training policy...")
    
    wandb.init(
        project=cfg.policy,
        group=cfg.task.task_name,
        name=f"v1.2.1",
        # track hyperparameters and run metadata
        config={
            "policy": cfg.policy,
            "dataset_type": cfg.task.dataset_type,
            "n_epochs": cfg.num_epochs,
            "seed": cfg.seed,
            "task": cfg.task.task_name,
        },
        # mode="disabled",
    )
    # wandb.watch(policy.model, log="all")
    # override Hydra parameters with WandB parameters
    cfg.policy.lr = wandb.config.lr
    cfg.policy.denoising_network.hidden_dim = wandb.config.hidden_dim
    cfg.policy.denoising_network.num_layers = wandb.config.num_layers
    cfg.policy.ckpt_path = cfg.policy.ckpt_path + f"_lr{cfg.policy.lr}_hd{str(wandb.config.hidden_dim)}_nl{str(wandb.config.num_layers)}"
    cfg.task.control_mode = wandb.config.control_mode

    # instanciate policy from cfg file
    policy = hydra.utils.instantiate(cfg.policy)
    log.info(f"Training policy {policy.__class__.__name__} with seed {cfg.seed} on task {cfg.task.task_name}")
    try:
        if cfg.policy.ckpt_path is not None and cfg.load_ckpt:
            policy.load_nets(cfg.policy.ckpt_path)
    except Exception as e:
        log.error(f"Error loading checkpoint: {e}")
    

    
    # Split the dataset into train and validation
    train_dataset, val_dataset = torch.utils.data.random_split(
        policy.dataset, [len(policy.dataset) - int(cfg.val_fraction * len(policy.dataset)), int(cfg.val_fraction * len(policy.dataset))]
    )
    

    E = cfg.num_epochs
    V = cfg.num_epochs
    if cfg.eval_params != "disabled":
        E = cfg.eval_params.eval_every
        V = cfg.eval_params.val_every
        assert V <= E, "Validation interval should be smaller than evaluation interval"
        assert E % V == 0, "Evaluation interval should be multiple of validation interval"
    
    try:
        policy.num_epochs = cfg.num_epochs
    except:
        log.error("Error setting total num_epochs in policy")
    
    # evaluate every E epochs
    max_success_rate = 0
    for i in range(1, 1 + cfg.num_epochs // V):
        # train policy
        policy.train(dataset=train_dataset,
                    num_epochs=V,
                    model_path=cfg.policy.ckpt_path,
                    seed=cfg.seed)
        log.info(f"Calculating validation loss...")
        val_loss = policy.validate(
                dataset=val_dataset,
                model_path=cfg.policy.ckpt_path,
            )
        wandb.log({"validation_loss": val_loss})
        # evaluate policy
        if i % (E/V) == 0:
            success_rate = eval_main(cfg.eval_params)
            if success_rate >= max_success_rate:
                max_success_rate = success_rate
                policy.save_nets(cfg.policy.ckpt_path + f"_best_succ={success_rate}.pt")


    wandb.finish()

if __name__ == "__main__":
    wandb.login()


    sweep_configuration = {
        "method": "random",
        "metric": {"name": "validation_loss", "goal": "minimize"},
        "parameters": {
            "lr": {"values": [0.0001, 0.00005, 0.00001]},
            "hidden_dim": {"values": [128, 256, 512, 1024]},
            "num_layers": {"values": [2, 3, 4, 5]},
            "control_mode": {"values": ["JOINT_VELOCITY"]}
        },
    }
    # 3: Start the sweep
    # sweep_id = wandb.sweep(sweep=sweep_configuration, project="GraphDDPM-Sweep")
    sweep_id = "8m873pjt"
    wandb.agent(sweep_id, function=train, count=20, project="GraphDDPM-Sweep")