imitate_episodes_pmi.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
import os
import pickle
import argparse
import matplotlib.pyplot as plt
from copy import deepcopy
from tqdm import tqdm
from einops import rearrange
from torchvision.models import resnet18, ResNet18_Weights

from constants import DT
from constants import PUPPET_GRIPPER_JOINT_OPEN
from utils import load_data_MINE # data functions
from utils import sample_box_pose, sample_insertion_pose # robot functions
from utils import compute_dict_mean, set_seed, detach_dict # helper functions
from policy import ACTPolicy_PMI, CNNMLPPolicy, Predict_Model, VAEPredictionModel, MINEEstimator
from visualize_episodes import save_videos

from detr.models.transformer import Transformer
from detr.models.backbone import build_backbone


import wandb
from sim_env import BOX_POSE

import IPython
e = IPython.embed
import subprocess
def get_free_gpu():
    try:
        output = subprocess.check_output(['nvidia-smi', '--query-gpu=memory.used,memory.total', '--format=csv,nounits,noheader'])
        lines = output.decode('utf-8').strip().split('\n')
        gpu_memory = [list(map(int, line.split(','))) for line in lines]
        free_memory_percentage = [(total - used) / total * 100 for used, total in gpu_memory]
        most_free_gpu = free_memory_percentage.index(max(free_memory_percentage))
        print(f"Selected GPU {most_free_gpu} with {free_memory_percentage[most_free_gpu]:.2f}% free memory")
        return most_free_gpu
    except Exception as e:
        print(f"Error selecting GPU: {e}")
        return 0  # Default to GPU 0 if there's an error

def main(args):
    # set_seed(1)
    seed = args['seed']
    set_seed(seed)
    # command line parameters
    is_eval = args['eval']
    ckpt_dir = args['ckpt_dir']
    policy_class = args['policy_class']
    onscreen_render = args['onscreen_render']
    task_name = args['task_name']
    batch_size_train = args['batch_size']
    batch_size_val = args['batch_size']
    mine_batch_size = args['mine_batch_size']
    num_epochs = args['num_epochs']
    load_mine = args['load_mine']

    # get task parameters
    is_sim = task_name[:4] == 'sim_'
    if is_sim:
        from constants import SIM_TASK_CONFIGS
        task_config = SIM_TASK_CONFIGS[task_name]
    else:
        from aloha_scripts.constants import TASK_CONFIGS
        task_config = TASK_CONFIGS[task_name]
    # print(is_sim)
    # print(task_config)
    # exit(0)
    dataset_dir = task_config['dataset_dir']
    num_episodes = task_config['num_episodes']
    episode_len = task_config['episode_len']
    camera_names = task_config['camera_names']

    # fixed parameters
    state_dim = 14
    lr_backbone = 1e-5
    backbone = 'resnet18'
    if policy_class == 'ACT':
        enc_layers = 4
        dec_layers = 7
        nheads = 8
        policy_config = {'lr': args['lr'],
                         'num_queries': args['chunk_size'],
                         'query_freq': args['query_freq'],
                         'kl_weight': args['kl_weight'],
                         'decay_rate': args['decay_rate'],
                         'hidden_dim': args['hidden_dim'],
                         'dim_feedforward': args['dim_feedforward'],
                         'lr_backbone': lr_backbone,
                         'backbone': backbone,
                         'enc_layers': enc_layers,
                         'dec_layers': dec_layers,
                         'nheads': nheads,
                         'camera_names': camera_names,
                         'state_dim': args['state_dim'],
                         'action_dim': args['action_dim'],
                         'reweight': args['reweight'],
                         'self_normalize': args['self_normalize_weight'],
                         'weight_clip': args['weight_clip'],
                         }
    elif policy_class == 'CNNMLP':
        policy_config = {'lr': args['lr'], 'lr_backbone': lr_backbone, 'backbone' : backbone, 'num_queries': 1,
                         'camera_names': camera_names,}
    else:
        raise NotImplementedError

    ckpt_dir = args['ckpt_dir'] + '_seed_' + str(seed)
    config = {
        'num_epochs': args['num_epochs_mine'],
        'ckpt_dir': ckpt_dir,
        'episode_len': episode_len,
        'state_dim': state_dim,
        'lr': args['lr'],
        'policy_class': policy_class,
        'onscreen_render': onscreen_render,
        'policy_config': policy_config,
        'task_name': task_name,
        'seed': args['seed'],
        'temporal_agg': args['temporal_agg'],
        'camera_names': camera_names,
        'real_robot': not is_sim,
        'use_predict_model': args['use_predict_model'],
    }

    if is_eval:
        print("EVAL MODE")
        ckpt_names = [f'policy_best.ckpt']
        results = []
        for ckpt_name in ckpt_names:
            success_rate, avg_return = eval_bc(config, ckpt_name, save_episode=True)
            results.append([ckpt_name, success_rate, avg_return])

        for ckpt_name, success_rate, avg_return in results:
            print(f'{ckpt_name}: {success_rate=} {avg_return=}')
        print()
        return

    train_dataloader, val_dataloader, train_dataloader_mine, val_dataloader_mine, stats, _ = load_data_MINE(dataset_dir, num_episodes, camera_names, batch_size_train, batch_size_val, mine_batch_size, args['chunk_size'])

    # save dataset stats
    if not os.path.isdir(ckpt_dir):
        os.makedirs(ckpt_dir)
    stats_path = os.path.join(ckpt_dir, f'dataset_stats.pkl')
    with open(stats_path, 'wb') as f:
        pickle.dump(stats, f)

    mine_config = config.copy()
    # prediction_config['policy_config']['action_dim'] = 512
    mine_ckpt_dir = mine_config['ckpt_dir'] = args['mine_ckpt_dir'] + '_seed_' + str(seed)
    mine_config['lr'] = mine_config['policy_config']['lr'] = args['lr_mine']
    mine_ckpt_path = os.path.join(mine_ckpt_dir, f'mine_model_best.ckpt')
    if not os.path.isdir(mine_ckpt_dir):
        os.makedirs(mine_ckpt_dir)
    if not load_mine:
        best_ckpt_info = train_mine(train_dataloader_mine, val_dataloader_mine, mine_config)
        best_epoch, min_val_loss, best_state_dict = best_ckpt_info
        # save best checkpoint
        torch.save(best_state_dict, mine_ckpt_path)
        print(f'MINE Best ckpt, val loss {min_val_loss:.6f} @ epoch{best_epoch}')

    # return
    ###################load prediction model
    # # predict_model_dir = '/localdata/yy/zzzzworkspace/act/ckpt/sim_transfer_cube_scripted_run2_decay_1/prediction_best_model.ckpt'
    # predict_model_dir = '/localdata/yy/zzzzworkspace/act/ckpt/sim_transfer_cube_scripted_run2_decay_1/prediction_model_epoch_1900_seed_1.ckpt'
    # predict_model_dir = '/localdata/yy/zzzzworkspace/act/ckpt/sim_transfer_cube_scripted_run2_decay_1/prediction_model_epoch_1957_seed_1.ckpt'
    # config['policy_config']['action_dim'] = 512
    # predict_model = make_predict_model(mine_config['policy_config'])
    # predict_model.load_state_dict(best_state_dict)
    # config['policy_config']['action_dim'] = 14
    # predict_model.cuda()
    mine_model = make_mine_estimator(mine_config['policy_config'])
    if not load_mine:
        mine_model.load_state_dict(best_state_dict)
    else:
        load_mine_ckpt_path = os.path.join(args['load_mine_ckpt_path'])
        mine_model.load_state_dict(torch.load(load_mine_ckpt_path))
    mine_model.cuda()

    config['policy_config']['action_dim'] = 14
    
    # train policy
    config['num_epochs'] = args['num_epochs']
    # config['action_dim'] = 14

    best_ckpt_info = train_bc(train_dataloader, val_dataloader, config, mine_model)
    best_epoch, min_val_loss, best_state_dict = best_ckpt_info

    # save best checkpoint
    ckpt_path = os.path.join(ckpt_dir, f'policy_best.ckpt')
    torch.save(best_state_dict, ckpt_path)
    print(f'Best ckpt, val loss {min_val_loss:.6f} @ epoch{best_epoch}')


    ### eval best ckpt
    ckpt_names = [f'policy_best.ckpt']
    results = []
    for ckpt_name in ckpt_names:
        success_rate, avg_return = eval_bc(config, ckpt_name, save_episode=True)
        results.append([ckpt_name, success_rate, avg_return])
        wandb.log({"best/success_rate": success_rate})
        wandb.log({"best/avg_return": avg_return})

    for ckpt_name, success_rate, avg_return in results:
        print(f'{ckpt_name}: {success_rate=} {avg_return=}')
    print()
    exit()


def make_policy(policy_class, policy_config):
    print(policy_config, "policy_config")
    # exit(0)
    if policy_class == 'ACT':
        policy = ACTPolicy_PMI(policy_config)
    elif policy_class == 'CNNMLP':
        policy = CNNMLPPolicy(policy_config)
    else:
        raise NotImplementedError
    # exit(0)
    return policy

def make_predict_model(policy_config):
    model = Predict_Model(policy_config)
    # model = VAEPredictionModel(policy_config)
    return model

def make_mine_estimator(policy_config):
    model = MINEEstimator(policy_config)
    return model

def make_optimizer(policy_class, policy):
    if policy_class == 'ACT':
        optimizer = policy.configure_optimizers()
    elif policy_class == 'CNNMLP':
        optimizer = policy.configure_optimizers()
    else:
        raise NotImplementedError
    return optimizer


def get_image(ts, camera_names):
    curr_images = []
    for cam_name in camera_names:
        curr_image = rearrange(ts.observation['images'][cam_name], 'h w c -> c h w')
        curr_images.append(curr_image)
    curr_image = np.stack(curr_images, axis=0)
    curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0)
    return curr_image


def eval_bc(config, ckpt_name, save_episode=True):
    set_seed(1000)
    print("EVAL BC-----------------------------------\n")
    ckpt_dir = config['ckpt_dir']
    state_dim = config['state_dim']
    real_robot = config['real_robot']
    policy_class = config['policy_class']
    onscreen_render = config['onscreen_render']
    policy_config = config['policy_config']
    camera_names = config['camera_names']
    max_timesteps = config['episode_len']
    task_name = config['task_name']
    temporal_agg = config['temporal_agg']
    onscreen_cam = 'angle'

    # load policy and stats
    ckpt_path = os.path.join(ckpt_dir, ckpt_name)
    policy = make_policy(policy_class, policy_config)

    loading_status = policy.load_state_dict(torch.load(ckpt_path))
    print(loading_status)
    policy.cuda()
    policy.eval()
    print(f'Loaded: {ckpt_path}')
    stats_path = os.path.join(ckpt_dir, f'dataset_stats.pkl')
    with open(stats_path, 'rb') as f:
        stats = pickle.load(f)

    pre_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std']
    post_process = lambda a: a * stats['action_std'] + stats['action_mean']

    # load environment
    if real_robot:
        from aloha_scripts.robot_utils import move_grippers # requires aloha
        from aloha_scripts.real_env import make_real_env # requires aloha
        env = make_real_env(init_node=True)
        env_max_reward = 0
    else:
        from sim_env import make_sim_env
        env = make_sim_env(task_name)
        env_max_reward = env.task.max_reward

    query_frequency = policy_config['num_queries']
    if temporal_agg:
        query_frequency = 1
        num_queries = policy_config['num_queries']
    else:
        query_frequency = policy_config['query_freq']
    ## yuyue add param here


    max_timesteps = int(max_timesteps * 1) # may increase for real-world tasks

    num_rollouts = 50
    episode_returns = []
    highest_rewards = []
    for rollout_id in range(num_rollouts):
        rollout_id += 0
        ### set task
        if 'sim_transfer_cube' in task_name:
            BOX_POSE[0] = sample_box_pose() # used in sim reset
        elif 'sim_insertion' in task_name:
            BOX_POSE[0] = np.concatenate(sample_insertion_pose()) # used in sim reset

        ts = env.reset()

        ### onscreen render
        if onscreen_render:
            ax = plt.subplot()
            plt_img = ax.imshow(env._physics.render(height=480, width=640, camera_id=onscreen_cam))
            plt.ion()

        ### evaluation loop
        if temporal_agg:
            all_time_actions = torch.zeros([max_timesteps, max_timesteps+num_queries, state_dim]).cuda()

        qpos_history = torch.zeros((1, max_timesteps, state_dim)).cuda()
        image_list = [] # for visualization
        qpos_list = []
        target_qpos_list = []
        rewards = []
        with torch.inference_mode():
            for t in range(max_timesteps):
                ### update onscreen render and wait for DT
                if onscreen_render:
                    image = env._physics.render(height=480, width=640, camera_id=onscreen_cam)
                    plt_img.set_data(image)
                    plt.pause(DT)

                ### process previous timestep to get qpos and image_list
                obs = ts.observation
                if 'images' in obs:
                    image_list.append(obs['images'])
                else:
                    image_list.append({'main': obs['image']})
                qpos_numpy = np.array(obs['qpos'])
                qpos = pre_process(qpos_numpy)
                qpos = torch.from_numpy(qpos).float().cuda().unsqueeze(0)
                qpos_history[:, t] = qpos
                curr_image = get_image(ts, camera_names)

                ### query policy
                if config['policy_class'] == "ACT":
                    if t % query_frequency == 0:
                        all_actions = policy(qpos, curr_image)
                        # print(all_actions.shape, "all_actions")
                        # exit(0)
                    if temporal_agg:
                        all_time_actions[[t], t:t+num_queries] = all_actions
                        actions_for_curr_step = all_time_actions[:, t]
                        actions_populated = torch.all(actions_for_curr_step != 0, axis=1)
                        actions_for_curr_step = actions_for_curr_step[actions_populated]
                        k = 0.01
                        exp_weights = np.exp(-k * np.arange(len(actions_for_curr_step)))
                        exp_weights = exp_weights / exp_weights.sum()
                        exp_weights = torch.from_numpy(exp_weights).cuda().unsqueeze(dim=1)
                        raw_action = (actions_for_curr_step * exp_weights).sum(dim=0, keepdim=True)
                    else:
                        raw_action = all_actions[:, t % query_frequency]
                elif config['policy_class'] == "CNNMLP":
                    raw_action = policy(qpos, curr_image)
                else:
                    raise NotImplementedError

                ### post-process actions
                raw_action = raw_action.squeeze(0).cpu().numpy()
                action = post_process(raw_action)
                target_qpos = action

                ### step the environment
                ts = env.step(target_qpos)

                ### for visualization
                qpos_list.append(qpos_numpy)
                target_qpos_list.append(target_qpos)
                rewards.append(ts.reward)

            plt.close()
        if real_robot:
            move_grippers([env.puppet_bot_left, env.puppet_bot_right], [PUPPET_GRIPPER_JOINT_OPEN] * 2, move_time=0.5)  # open
            pass

        rewards = np.array(rewards)
        episode_return = np.sum(rewards[rewards!=None])
        episode_returns.append(episode_return)
        episode_highest_reward = np.max(rewards)
        highest_rewards.append(episode_highest_reward)
        print(f'Rollout {rollout_id}\n{episode_return=}, {episode_highest_reward=}, {env_max_reward=}, Success: {episode_highest_reward==env_max_reward}')
        wandb.log({"rollout/episode_return": episode_return})
        wandb.log({"rollout/episode_highest_reward": episode_highest_reward})
        wandb.log({"rollout/success": episode_highest_reward==env_max_reward})
        if save_episode:
            save_videos(image_list, DT, video_path=os.path.join(ckpt_dir, f'video{rollout_id}.mp4'),gif_path=os.path.join(ckpt_dir, f'video{rollout_id}.gif'))
            # wandb.log({"rollout/video": wandb.Video(os.path.join(ckpt_dir, f'video{rollout_id}.mp4'))})
            # wandb.log({"rollout/gif": wandb.Video(os.path.join(ckpt_dir, f'video{rollout_id}.gif'))})

    success_rate = np.mean(np.array(highest_rewards) == env_max_reward)
    avg_return = np.mean(episode_returns)
    wandb.log({"rollout/success_rate": success_rate})
    wandb.log({"rollout/avg_return": avg_return})
    summary_str = f'\nSuccess rate: {success_rate}\nAverage return: {avg_return}\n\n'
    for r in range(env_max_reward+1):
        more_or_equal_r = (np.array(highest_rewards) >= r).sum()
        more_or_equal_r_rate = more_or_equal_r / num_rollouts
        summary_str += f'Reward >= {r}: {more_or_equal_r}/{num_rollouts} = {more_or_equal_r_rate*100}%\n'

    print(summary_str)

    # save success rate to txt
    result_file_name = 'result_' + ckpt_name.split('.')[0] + '_query_freq_' + str(policy_config['query_freq']) + '_temporal_agg_' + str(temporal_agg) + '.txt'
    with open(os.path.join(ckpt_dir, result_file_name), 'w') as f:
        f.write(summary_str)
        f.write(repr(episode_returns))
        f.write('\n\n')
        f.write(repr(highest_rewards))

    return success_rate, avg_return


def forward_pass_vanilla(data, policy):
    image_data, qpos_data, action_data, is_pad = data
    image_data, qpos_data, action_data, is_pad = image_data.cuda(), qpos_data.cuda(), action_data.cuda(), is_pad.cuda()
    return policy(qpos_data, image_data, action_data, is_pad) # TODO remove None

def forward_pass_with_prediction(data, policy, model):
    image_data, image_features, qpos_data, action_data, is_pad = data
    image_data, qpos_data, image_features, action_data, is_pad = image_data.cuda(), qpos_data.cuda(), image_features.cuda(), action_data.cuda(), is_pad.cuda()
    return policy(qpos_data, image_data, action_data, is_pad, model,image_features)
# def eval_from_policy(policy):

def forward_pass_with_mine(data, policy, model):
    image_data, qpos_data, action_data, is_pad = data
    image_data, qpos_data, action_data, is_pad = image_data.cuda(), qpos_data.cuda(), action_data.cuda(), is_pad.cuda()
    return policy(qpos_data, image_data, action_data, is_pad, model)


def train_bc(train_dataloader, val_dataloader, config, mine_model):
    num_epochs = config['num_epochs']
    ckpt_dir = config['ckpt_dir']
    seed = config['seed']
    policy_class = config['policy_class']
    policy_config = config['policy_config']
    use_predict_model = config['use_predict_model']
    set_seed(seed)

    policy = make_policy(policy_class, policy_config)
    policy.cuda()
    optimizer = make_optimizer(policy_class, policy)

    train_history = []
    validation_history = []
    min_val_loss = np.inf
    best_ckpt_info = None
    for epoch in tqdm(range(num_epochs)):
        print(f'\nEpoch {epoch}')
        # validation
        with torch.inference_mode():
            policy.eval()
            epoch_dicts = []
            for batch_idx, data in enumerate(val_dataloader):
                forward_dict = forward_pass_vanilla(data, policy)
                epoch_dicts.append(forward_dict)
            epoch_summary = compute_dict_mean(epoch_dicts)
            validation_history.append(epoch_summary)

            epoch_val_loss = epoch_summary['loss']
            if epoch_val_loss < min_val_loss:
                min_val_loss = epoch_val_loss
                best_ckpt_info = (epoch, min_val_loss, deepcopy(policy.state_dict()))
        print(f'Val loss:   {epoch_val_loss:.5f}')
        wandb.log({"val_loss": epoch_val_loss})
        summary_string = ''
        for k, v in epoch_summary.items():
            summary_string += f'{k}: {v.item():.3f} '
            wandb.log({k: v.item()})
        print(summary_string)

        # training
        policy.train()
        optimizer.zero_grad()
        for batch_idx, data in enumerate(train_dataloader):
            forward_dict = forward_pass_with_mine(data, policy, mine_model)
            # backward
            loss = forward_dict['loss']
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            train_history.append(detach_dict(forward_dict))
        epoch_summary = compute_dict_mean(train_history[(batch_idx+1)*epoch:(batch_idx+1)*(epoch+1)])
        epoch_train_loss = epoch_summary['loss']
        print(f'Train loss: {epoch_train_loss:.5f}')
        wandb.log({"train_loss": epoch_train_loss})
        summary_string = ''
        for k, v in epoch_summary.items():
            summary_string += f'{k}: {v.item():.3f} '
        print(summary_string)

        if epoch % 200 == 0:
            ckpt_path = os.path.join(ckpt_dir, f'policy_epoch_{epoch}_seed_{seed}.ckpt')
            torch.save(policy.state_dict(), ckpt_path)
            plot_history(train_history, validation_history, epoch, ckpt_dir, seed)

    ckpt_path = os.path.join(ckpt_dir, f'policy_last.ckpt')
    torch.save(policy.state_dict(), ckpt_path)

    best_epoch, min_val_loss, best_state_dict = best_ckpt_info
    ckpt_path = os.path.join(ckpt_dir, f'policy_epoch_{best_epoch}_seed_{seed}.ckpt')
    torch.save(best_state_dict, ckpt_path)
    print(f'Training finished:\nSeed {seed}, val loss {min_val_loss:.6f} at epoch {best_epoch}')

    # save training curves
    plot_history(train_history, validation_history, num_epochs, ckpt_dir, seed)

    return best_ckpt_info


def plot_history(train_history, validation_history, num_epochs, ckpt_dir, seed):
    # save training curves
    for key in train_history[0]:
        plot_path = os.path.join(ckpt_dir, f'train_val_{key}_seed_{seed}.png')
        plt.figure()
        train_values = [summary[key].item() for summary in train_history]
        val_values = [summary[key].item() for summary in validation_history]
        plt.plot(np.linspace(0, num_epochs-1, len(train_history)), train_values, label='train')
        plt.plot(np.linspace(0, num_epochs-1, len(validation_history)), val_values, label='validation')
        # plt.ylim([-0.1, 1])
        plt.tight_layout()
        plt.legend()
        plt.title(key)
        plt.savefig(plot_path)
    print(f'Saved plots to {ckpt_dir}')

def train_mine(train_dataloader, val_dataloader, config):
    num_epochs = config['num_epochs']
    ckpt_dir = config['ckpt_dir']
    seed = config['seed']
    # policy_class = config['policy_class']
    policy_config = config['policy_config']

    set_seed(seed)

    model = make_mine_estimator(policy_config)
    model.cuda()
    optimizer = model.configure_optimizers()

    min_val_loss = np.inf
    best_ckpt_info = None
    train_history = []
    validation_history = []
    for epoch in tqdm(range(num_epochs)):
        # validation
        with torch.inference_mode():
            model.eval()
            epoch_loss_list = []
            epoch_total_loss_list = []
            pos_exp_list = []
            neg_exp_list = []
            for batch_idx, data in enumerate(val_dataloader):
                # image_data, image_rep, qpos_data, action_data, is_pad = data
                # image_data, image_rep, qpos_data, action_data, is_pad = image_data.cuda(), image_rep.cuda(), qpos_data.cuda(), action_data.cuda(), is_pad.cuda()
                # _, L = model(qpos_data, image_data, image_rep, is_pad, is_training=False)
                # epoch_dicts.append(L)
                image_data, qpos_data, action_data, timesteps = data
                image_data, qpos_data, action_data, timesteps = image_data.cuda(), qpos_data.cuda(), action_data.cuda(), timesteps.cuda()
                loss_dict, pos, neg = model(image_data, action_data, timesteps, qpos_data)
                epoch_loss_list.append(loss_dict['loss'].item())
                epoch_total_loss_list.append(loss_dict['total_loss'].item())
                pos_exp_list.append(pos.mean().item())
                neg_exp_list.append(neg.mean().item())
            epoch_summary = {'loss': np.mean(epoch_loss_list), 'total_loss': np.mean(epoch_total_loss_list), 'pos_exp': np.mean(pos_exp_list), 'neg_exp': np.mean(neg_exp_list)}

            epoch_val_loss = epoch_summary['loss']
            if epoch_val_loss < min_val_loss:
                min_val_loss = epoch_val_loss
                best_ckpt_info = (epoch, min_val_loss, deepcopy(model.state_dict()))
        print(f'Val loss:   {epoch_val_loss:.10f}')
        wandb.log({"MINE/val_loss": epoch_val_loss})
        summary_string = ''
        for k, v in epoch_summary.items():
            summary_string += f'{k}: {v.item():.10f} '
            wandb.log({f"MINE/{k}": v.item()})


        optimizer.zero_grad()
        for batch_idx, data in enumerate(train_dataloader):
            image_data, qpos_data, action_data, timesteps = data
            image_data, qpos_data, action_data, timesteps = image_data.cuda(), qpos_data.cuda(), action_data.cuda(), timesteps.cuda()
            # loss.backward()
            # print(image_data.shape, image_rep.shape, qpos_data.shape, action_data.shape, is_pad.shape,end=" shapes of image_data, image_rep, qpos_data, action_data, is_pad\n")
            loss_dict, _, _ = model(image_data, action_data, timesteps, qpos_data)
            # loss = L['loss']
            # print(loss.item(), "loss\n")
            loss = loss_dict['total_loss']
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            train_history.append(loss.item())
        epoch_summary = {'loss': np.mean(train_history[(batch_idx+1)*epoch:(batch_idx+1)*(epoch+1)])} # this is total loss
        epoch_train_loss = epoch_summary['loss']
        print(f'Train loss: {epoch_train_loss:.10f}')
        wandb.log({"MINE/train_loss": epoch_train_loss})
        summary_string = ''
        for k, v in epoch_summary.items():
            summary_string += f'{k}: {v.item():.10f} '
            wandb.log({f"MINE/{k}": v.item()})
        print(summary_string)
        # Save the model after each epoch
        if epoch % 100 == 0:
            ckpt_path = os.path.join(ckpt_dir, f'mine_model_epoch_{epoch}_seed_{seed}.ckpt')
            torch.save(model.state_dict(), ckpt_path)

    ckpt_path = os.path.join(ckpt_dir, f'mine_model_last.ckpt')
    torch.save(model.state_dict(), ckpt_path)

    best_epoch, min_val_loss, best_state_dict = best_ckpt_info
    ckpt_path = os.path.join(ckpt_dir, f'best_mine_model_epoch_{best_epoch}_seed_{seed}.ckpt')
    torch.save(best_state_dict, ckpt_path)
    print(f'Training finished:\nSeed {seed}, val loss {min_val_loss:.6f} at epoch {best_epoch}')

    return best_ckpt_info

if __name__ == '__main__':
    wandb.init()
    config = wandb.config
    config['seed'] = 42
    selected_gpu = get_free_gpu()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_gpu)
    # parser = argparse.ArgumentParser()
    # parser.add_argument('--eval', action='store_true')
    # parser.add_argument('--onscreen_render', action='store_true')
    # parser.add_argument('--ckpt_dir', action='store', type=str, help='ckpt_dir', required=True)
    # parser.add_argument('--policy_class', action='store', type=str, help='policy_class, capitalize', required=True)
    # parser.add_argument('--task_name', action='store', type=str, help='task_name', required=True)
    # parser.add_argument('--batch_size', action='store', type=int, help='batch_size', required=True)
    # parser.add_argument('--seed', action='store', type=int, help='seed', required=True)
    # parser.add_argument('--num_epochs', action='store', type=int, help='num_epochs', required=True)
    # parser.add_argument('--lr', action='store', type=float, help='lr', required=True)

    # # for ACT
    # parser.add_argument('--kl_weight', action='store', type=int, help='KL Weight', required=False)
    # parser.add_argument('--chunk_size', action='store', type=int, help='chunk_size', required=False)
    # parser.add_argument('--hidden_dim', action='store', type=int, help='hidden_dim', required=False)
    # parser.add_argument('--dim_feedforward', action='store', type=int, help='dim_feedforward', required=False)
    # parser.add_argument('--temporal_agg', action='store_true')
    # print(vars(parser.parse_args()))
    # print(config)
    # exit(0)
    # main(vars(parser.parse_args()))
    main(config)