train.py

import torch
import sys
import os
import numpy as np
import argparse
import time

import matplotlib.pyplot as plt

import wandb

import sys
sys.path.append("./CleanDiffuser")

from dsp.dataset.surrol_dataset import SurrolDataset, MixedSurrolDataset
from dsp.model.diffusion_loader import build_diffusion_agent
from dsp.utils.filter_thresholds import calc_kmeans_threshold, calc_logGMM_threshold, calc_otsu_threshold
from dsp.utils.metrics import compute_binary_classification_metrics
from CleanDiffuser.cleandiffuser.utils import loop_dataloader
from torch.utils.data import DataLoader

def train(args):
    # wandb parameters
    if args.msg != '':
        msg = '_' + args.msg
    project_name = 'surrol'
    run_name = 'diffusion_' + args.env + '_' + args.msg
    if args.isLog:
        wandb.init(project=project_name, name=run_name)
    
    ############################################
    # shape:                                   #
    # |O||O||O||O|| || || |                    #
    # | || || ||A||A||A||A|                    #
    # horizon = 7                              #
    # pad_before = 3                           #
    # pad_after = 3                            #
    ############################################
    # shape:                                   #
    # |O|                                      #
    # |A|                                      #
    # horizon = 1                              #
    # pad_before = 0                           #
    # pad_after = 0                            #
    ############################################

    To = int(args.oaPair[0])
    Ta = int(args.oaPair[2])
    horizon = To+Ta-1
    pad_before = To-1
    pad_after = Ta-1

    device = "cuda" if torch.cuda.is_available() else "cpu"

    ########################################################################################
    ## set up dataset
    ########################################################################################

    if args.clean_dataset_name != '' and args.noisy_dataset_name != '':
        # when training with mixed dataset, we need to calculate the normalizer using all data
        print("\nuse mixed dataset")
        mixed_dataset = MixedSurrolDataset(
            dataset_root=args.dataset_root,
            clean_dataset_path=args.clean_dataset_name,
            noisy_dataset_path=args.noisy_dataset_name,
            horizon=horizon, pad_before=pad_before, pad_after=pad_after,
            task_name=args.env, device=device
        )
        normalizer = mixed_dataset.get_normalizer()

        if args.train_for_clean:
            # stage 1
            print("training for stage1")
            print(f"Loaded Clean dataset from {args.clean_dataset_name}.npz")
            dataset = SurrolDataset(dataset_root=args.dataset_root,
                                    dataset_path=args.clean_dataset_name,
                                    horizon=horizon, pad_before=pad_before, pad_after=pad_after,
                                    task_name=args.env, device=device, normalizer=normalizer)
        else:
            # stage 2
            print("training for stage2")
            print(f"Loaded Clean dataset from {args.clean_dataset_name}.npz")
            clean_dataset = SurrolDataset(dataset_root=args.dataset_root,
                dataset_path=args.clean_dataset_name,
                horizon=horizon, pad_before=pad_before, pad_after=pad_after,
                task_name=args.env, device=device, normalizer=normalizer)
            print(f"Loaded Noisy dataset from {args.noisy_dataset_name}.npz")
            noisy_dataset = SurrolDataset(dataset_root=args.dataset_root,
                dataset_path=args.noisy_dataset_name,
                horizon=horizon, pad_before=pad_before, pad_after=pad_after,
                task_name=args.env, device=device, normalizer=normalizer, noisy_dataset=True)
            dataset = clean_dataset+noisy_dataset
    else:
        # train with single dataset
        print("use single dataset")
        dataset_name = args.clean_dataset_name if args.clean_dataset_name != '' else args.noisy_dataset_name
        dataset = SurrolDataset(dataset_root=args.dataset_root,
            dataset_path=dataset_name,
            horizon=horizon, pad_before=pad_before, pad_after=pad_after,
            task_name=args.env, device=device)
    
    data = dataset[0]
    obs, act = data["obs"]["state"], data["action"]
    obs_dim, act_dim = obs.shape[-1], act.shape[-1]
    act_dim = act_dim * (pad_after + 1)
    print(f'\nFinish loading data. Observation shape: {obs.shape}. Action shape {act.shape}.')



    ########################################################################################
    ## set up training mode
    ########################################################################################

    # get savename
    if args.debug:
        savename = time.strftime("%Y%m%d-%H%M%S") + "_" + (args.env+"-diffusion"+'_'+args.oaPair + msg + ".pt")
    else:
        savename = (args.env + "-diffusion" + '_' + args.oaPair + msg + ".pt")

    ## check if the model already exists
    if os.path.exists(args.savepath + '/' + savename) and not args.test_filter:
        print(f"\n###### MODEL EXIST, DID NOT TRAIN ######  model already exist! path: {args.savepath + '/' + savename}\n")
        return
    
    # setting up training mode according to training config
    print(f"\ntraining for task: {args.env}")
    
    ## whether to save the model
    if not args.no_save:
        print(f"\n####### SAVE MODEL ######## model will be saved at: {args.savepath + '/' + savename}")  
    else:
        print("\n####### NOT SAVED ########: no_save option is active, the model will not be saved!!!")
    
    ## whether to test the filter and get metrics
    if args.test_filter:
        print("\n####### TEST FILTER ########: test_filter option is active\n")
        os.makedirs(args.log_folder) if not os.path.exists(args.log_folder) else None 
        
        # threshold_type_count = 5
        # Accuracy_sum, Recall_sum, Precision_sum = [0.0] * threshold_type_count, [0.0] * threshold_type_count, [0.0] * threshold_type_count
        Accuracy_sum, Recall_sum, Precision_sum = 0.0, 0.0, 0.0

        if args.log_folder !='':
            log_output_file_path = os.path.join(args.log_folder, savename[:-3]+'.txt')
            print(f"logging filter test to {log_output_file_path}")
        else:
            raise ValueError('no log folder path provided!')

    ########################################################################################
    ## set up model
    ########################################################################################
    
    if args.checkpoint != '':
        actor = build_diffusion_agent(
            obs_dim=obs_dim,
            act_dim=act_dim,
            device=device,
            to_steps=To,
            model_path=os.path.join(args.savepath, args.checkpoint + '.pt'),
            mode="train",
        )
    else:
        actor = build_diffusion_agent(
            obs_dim=obs_dim,
            act_dim=act_dim,
            device=device,
            to_steps=To,
            mode="train",
        )
        print("did not load checkpoint!")

    ## whether to use online strategy
    if args.filter_noisy:
        print("\n######### FILTER MODE ########## filtering noise!")
        if args.online:
            print("\n####### ONLINE MODE ######## using online strategy!\n")
        else:
            print(f"\n####### OFFLINE MODE ######## loading independent judge:")
            filter_agent = build_diffusion_agent(
                obs_dim=obs_dim,
                act_dim=act_dim,
                device=device,
                to_steps=To,
                model_path=os.path.join(args.savepath, args.checkpoint + '.pt'),
                mode="eval",
            )
            
            



    ########################################################################################
    ## training model
    ########################################################################################
    
    ## start the training

    print(f"\ntraining for {args.train_step} step(s)!\n")

    if not os.path.exists(args.savepath):
        os.makedirs(args.savepath)

    ### setup dataloader
    dataloader = DataLoader(dataset, batch_size=256, shuffle=True, num_workers=3, persistent_workers=True)

    n_gradient_steps = 0
    avg_loss = 0.
    actor.train()

    #############################
    use_otsu_thr = False
    use_kmeans = False
    use_logGMM = False
    use_strict = False
    use_loose = False

    if args.custom_threshold == 'otsu':
        use_otsu_thr = True
    elif args.custom_threshold == 'k-means':
        use_kmeans = True
    elif args.custom_threshold == 'logGMM':
        use_logGMM = True
    elif args.custom_threshold == 'strict':
        use_strict = True
    elif args.custom_threshold == 'loose':
        use_loose = True
    ############################

    if args.custom_threshold != '':
        print("use custom thresholding method for filtering")

        if use_otsu_thr:
            print("use otsu to find the threshold for filtering")

        elif use_kmeans:
            print("use kmeans to find the threshold for filtering")

        elif use_logGMM:
            print("use logGMM to find the threshold for filtering")

        elif use_strict:
            print("use strict normal distribution thresholding (mean-var) for filtering")
        
        elif use_loose:
            print("use loose normal distribution thresholding (mean+var) for filtering")

        else:
            print("no custom thresholding method specified, use normal distribution thresholding (mean) for filtering")
    

    for batch in loop_dataloader(dataloader):

        obs, act = batch["obs"]["state"][:, :pad_before+1].to(device), batch["action"][:, pad_after:].to(device)
        act = act.reshape(act.size(0), -1)
        
        # TODO: currently, we only consider the situation of 1o1a. If we want to extend it to 4o4a, we need to
        #  pad the trajectory.

        ##################
        # filtering
        ##################
        if args.filter_noisy:
            with torch.no_grad():
                if not args.test_gt:
                    # generate data from diffusion

                    # for psm envs, we add noise to positions control
                    if args.env in {'NeedlePick-v0', 'NeedleReach-v0', 'PegTransfer-v0', 'GauzeRetrieve-v0'}:
                        act_mask = torch.tensor([[1, 1, 1, 0, 0]]).to(device)
                    elif args.env in {'BiPegTransfer-v0', 'NeedleRegrasp-v0'}:
                        act_mask = torch.tensor([[1, 1, 1, 0, 0, 1, 1, 1, 0, 0]]).to(device)
                    else:
                        raise NotImplementedError('Not implemented')
                    
                    # get filter mask from filter judge
                    if args.online:
                        sampled_acts, _ = actor.sample(
                            prior=torch.zeros_like(act), solver="ddpm", n_samples=obs.shape[0], sample_steps=5,
                            condition_cfg=obs, w_cfg=1.0)
                    else:
                        sampled_acts, _ = filter_agent.sample(
                            prior=torch.zeros_like(act), solver="ddpm", n_samples=obs.shape[0], sample_steps=5,
                            condition_cfg=obs, w_cfg=1.0)
                        
                    # calculate the threshold for filter
                    error = torch.sum(((sampled_acts - act) * act_mask) ** 2, dim=-1)
                    mean = torch.mean(error)
                    var = torch.var(error)

                    assert torch.all(error >= 0), "Error tensor contains negative values, which is not allowed."

                    ####################### check histogram ########################
                    # output_dir = 'output_images'
                    # os.makedirs(output_dir, exist_ok=True)

                    # error_np = error.detach().cpu().numpy()
                    # # Create log-spaced bins
                    # eps = 1e-8
                    # log_min = torch.log10(torch.min(error[error > 0]) + eps)
                    # log_max = torch.log10(torch.max(error) + eps)
                    # log_bins = torch.logspace(log_min, log_max, 40)  # 256 log-spaced bins

                    # # Plot histogram on normal scale
                    # plt.figure(figsize=(12, 6))
                    # plt.hist(error_np, bins=50, alpha=0.7, edgecolor='black')
                    # plt.xlabel('Error Value')
                    # plt.ylabel('Frequency')
                    # plt.title('Error Distribution Histogram')
                    # plt.grid(True, alpha=0.3)
                    # plt.savefig(os.path.join(output_dir, 'error_histogram.png'), dpi=150, bbox_inches='tight')
                    # plt.close()

                    # # Plot histogram
                    # plt.figure(figsize=(12, 6))
                    # plt.hist(error_np, bins=log_bins, alpha=0.7, edgecolor='black')
                    # plt.xscale('log')  # Keep log scale for proper visualization
                    # plt.xlabel('Error Value (log bins)')
                    # plt.ylabel('Frequency')
                    # plt.title('Histogram with Log-spaced Bins')
                    # plt.grid(True, alpha=0.3)
                    # plt.savefig(os.path.join(output_dir, 'error_histogram_log_bins.png'), dpi=150, bbox_inches='tight')
                    # plt.close()

                    ################################################################

                    if args.custom_threshold != '':

                        if use_otsu_thr:
                            # print('test: if using otsu')
                            thr = calc_otsu_threshold(error)
                            # print(f"otsu thr: {thr:.4f}, mean: {mean:.4f}, mean-var: {(mean-var):.4f}")

                        elif use_logGMM:
                            # print('test: if using logGMM')
                            thr = calc_logGMM_threshold(error)
                            # print(f"log-GMM thr: {thr:.4f}, mean: {mean:.4f}, mean-var: {(mean-var):.4f}")

                        elif use_kmeans:
                            # print('test: if using kmeans')
                            thr = calc_kmeans_threshold(error)
                            # print(f"k-means thr (original scale): {thr:.4f}, mean: {mean:.4f}, mean-var: {(mean-var):.4f}")

                        elif use_strict:
                            # print('test: if using strict')
                            thr = mean - var

                        elif use_loose:
                            thr = mean + var
                        
                        # for debug
                        
                        # print(f"mean: {mean:.4f}, mean-var: {(mean-var):.4f}")
                        # print(f"k-means: {kmeans_thr:.4f}, otsu: {otsu_thr:.4f}, gmm: {gmm_thr:.4f}")

                        # # Draw the histogram of error in log-scale bin and mark the thresholds
                        # plt.figure(figsize=(12, 6))
                        # plt.hist(error_np, bins=log_bins, alpha=0.7, edgecolor='black', label='Error Distribution')
                        # plt.xscale('log')  # Keep log scale for proper visualization

                        # # Mark the thresholds
                        # plt.axvline(mean.item(), color='blue', linestyle='--', label='Mean')
                        # plt.axvline((mean - var).item(), color='green', linestyle='--', label='Mean - Var')
                        # plt.axvline(kmeans_thr.item(), color='orange', linestyle='--', label='K-Means Threshold')
                        # plt.axvline(otsu_thr.item(), color='purple', linestyle='--', label='Otsu Threshold')
                        # plt.axvline(gmm_thr, color='red', linestyle='--', label='Log-GMM Threshold')

                        # # Add labels and legend
                        # plt.xlabel('Error Value (log bins)')
                        # plt.ylabel('Frequency')
                        # plt.title('Histogram with Log-spaced Bins and Thresholds')
                        # plt.legend()
                        # plt.grid(True, alpha=0.3)

                        # # Save the plot
                        # output_dir = 'output_images'
                        # os.makedirs(output_dir, exist_ok=True)
                        # plt.savefig(os.path.join(output_dir, 'error_histogram_with_thresholds.png'), dpi=150, bbox_inches='tight')
                        # plt.close()

                    else:
                        # print('test: if using normal')
                        thr = mean


                    noisy_idx = torch.logical_and((error > thr), batch['from_noisy_dataset'].squeeze(1).to(device))


                    if args.test_filter and n_gradient_steps % 10 == 0:
                        # test the filter
                        # print(n_gradient_steps)
                        gt_noisy_idx = batch['gt_noisy_idx'].squeeze(1)
                        # true_noisy_count = torch.sum(gt_noisy_idx).item()
                        # print(f"True amount of gt_noisy_idx: {true_noisy_count}")

                        # get confusion matrix(True: is noisy step)
                        Accuracy, Recall, Precision = compute_binary_classification_metrics(gt_noisy_idx.to(device), noisy_idx.to(device))
                        Accuracy_sum += Accuracy
                        Recall_sum += Recall 
                        Precision_sum += Precision
                                       
                else:
                    noisy_idx = batch['gt_noisy_idx'].squeeze(1)

                if torch.sum(~noisy_idx)<=1:
                    continue
                obs = obs[~noisy_idx]
                act = act[~noisy_idx]


        loss = actor.update(x0=act, condition=obs)["loss"]
        if args.isLog:
            wandb.log({"step_loss": loss})

        avg_loss += loss

        n_gradient_steps += 1

        if n_gradient_steps % 1000 == 0:
            print(f'Step: {n_gradient_steps} | Loss: {avg_loss / 1000}')

            if args.test_filter:
                test_message = f'step: {n_gradient_steps} | Accuracy: {(Accuracy_sum/100):.4f}, Recall: {(Recall_sum/100):.4f}, Precision: {(Precision_sum/100):.4f}\n'

                print(test_message)
                with open(log_output_file_path, "a") as file:
                    file.write(test_message)    

                Accuracy_sum = 0.0
                Recall_sum = 0.0
                Precision_sum = 0.0

                # Accuracy_sum = [0.0] * threshold_type_count
                # Recall_sum = [0.0] * threshold_type_count
                # Precision_sum = [0.0] * threshold_type_count

            if args.isLog:
                wandb.log({"step_loss": loss, "step": n_gradient_steps})
            avg_loss = 0.

        if (n_gradient_steps % args.save_step == 0) and not args.no_save:
            actor.save(args.savepath + "/" + f"step{n_gradient_steps/1000}k_"+savename)
            print(f"actor saved at step:{n_gradient_steps} \nlocated at: {args.savepath + '/' + savename}")

        if (n_gradient_steps == args.train_step) and not args.no_save:
            actor.save(args.savepath + "/" + savename)
            print(f"training complete. The final model is being saved at: {args.savepath + '/' + savename}.")
            break
        


if __name__ == '__main__':
    ##################### edit here ##########################
    """SurrolDataset is compatible with every dataset generated by data_generation.py in Surrol. 
    just edit the root is fine."""
    parser = argparse.ArgumentParser(description='args for training diffusion model')
    parser.add_argument('--env', type=str, required=True,
                        help='the environment to generate demonstrations')
    
    parser.add_argument('--dataset_root', type=str, default='datasets/demo',
                        help='root of the dataset')
    parser.add_argument('--clean_dataset_name', type=str, default='',
                        help='the name of the clean dataset')
    parser.add_argument('--noisy_dataset_name', type=str, default='',
                        help='the name of the noisy dataset')
    
    parser.add_argument('--checkpoint', type=str, default='',
                        help='whether to load model checkpoint')
    
    parser.add_argument('--train_for_clean', action='store_true',
                        help='train the diffusion model for clean dataset, using mixed normalizer')
    parser.add_argument('--filter_noisy', action='store_true',
                        help='filter the noisy data using trained diffusion model')
    parser.add_argument('--online', action='store_true', 
                        help='use online strategy to filter noisy data')
    parser.add_argument('--custom_threshold', type=str, choices=['otsu', 'k-means', 'logGMM', 'strict', 'loose'], default='',
                        help='use custom thresholding method(otsu or k-means) to filter noisy')

    
    parser.add_argument('--train_step', type=int, default=100_000,
                        help='number of training steps')
    parser.add_argument('--save_step', type=int, default=100_000,
                        help='save model every n steps')
    parser.add_argument('--savepath', type=str, default="trained_models",
                        help='path to save the trained model')
    parser.add_argument('--no_save', action='store_true',
                        help='do not save the trained model.')
    
    parser.add_argument('--test_gt', action='store_true',
                        help='use the ground-truth to filter the noisy data for benchmarking')
    parser.add_argument('--test_filter', action='store_true',
                        help='enable for testing filter without training')
    parser.add_argument('--log_folder', type=str, default='logs/filter_test',
                        help='folder to save logs')
    
    #TODO implement 4o4a
    parser.add_argument('--oaPair', type=str, default='1o1a', choices=['1o1a'],
                        help='use history obs, predict future action trunk')


    parser.add_argument('--debug', action='store_true',
                        help='debug mode')
    parser.add_argument('--isLog', action='store_true',
                        help='log to wandb')

    parser.add_argument('--msg', type=str, default='',
                        help='extra information for the run')
    args = parser.parse_args()

    print("\n##############################")
    print("##         training         ##")
    print("##############################\n")
    print(args)
    ##########################################################
    train(args)