eval.py

import torch
import gym
import os
import sys
import wandb
from tqdm import tqdm
import argparse
# from utils.envrionment import MatchBoardSLWrapper, BiPegBoardSLWrapper

# from components.logger import Logger, WandBLogger, logger
from dsp.utils.general_utils import set_seed_everywhere
from dsp.model.diffusion_loader import build_diffusion_agent
from dsp.modules.rollout_sampler import DiffusionRolloutSampler
from dsp.utils.video_utils import episode_to_mp4
from dsp.dataset.surrol_dataset import SurrolDataset, MixedSurrolDataset


class EvalSampler():
    def __init__(self, args):
        self.args = args
        self._device = "cuda:0" if torch.cuda.is_available() else "cpu"

        # environment
        self._env = gym.make(self.args.env, render_mode='rgb_array') # 'rgb_array', 'human'
        # if args.env=="MatchBoard-v0":
        #     print("using new code for MatchBoard-v0!")
        #     self._env = MatchBoardSLWrapper(self._env, output_raw_obs=True, subtask='pull')
        # if args.env=="BiPegBoard-v0":
        #     print("using new code for BiPegBoard-v0!")
        #     self._env = BiPegBoardSLWrapper(self._env, output_raw_obs=True, subtask='grasp')
        self._episode_step = 0
        self._max_episode_len = args.max_episode_length

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

        # diffusion agent
        self._load_diffusion_agent()
        self.rollout_sampler = DiffusionRolloutSampler(
            env=self._env,
            agent=self.agent,
            state_normalizer=self.state_normalizer,
            action_normalizer=self.action_normalizer,
            pad_before=self.pad_before,
            max_episode_len=self._max_episode_len,
            act_dim=self.act_dim,
            device=self._device,
        )

    def _load_diffusion_agent(self):
        # if args.env in ["MatchBoard-v0", "BiPegBoard-v0"]:
        #     obs_tmp, _ = self._env.reset()
        # else:
        #     obs_tmp = self._env.reset() 
        obs_tmp = self._env.reset()
        self.obs_dim = obs_tmp['observation'].shape[0] + obs_tmp['achieved_goal'].shape[0] + obs_tmp['desired_goal'].shape[0]
        self.act_dim = self._env.ACTION_SIZE * (self.pad_after + 1)
        # print("with env.reset() get:")
        # print(len(obs_tmp['observation']), len(obs_tmp['achieved_goal']), len(obs_tmp['desired_goal']))
        # print(self.obs_dim)

        model_path = os.path.join(self.args.model_root, self.args.model_name+'.pt')

        self.agent = build_diffusion_agent(
            obs_dim=self.obs_dim,
            act_dim=self.act_dim,
            device=self._device,
            to_steps=self.To,
            model_path=model_path,
            mode="eval",
        )

    def _load_normalizer_from_dataset(self):
        # load normalizer from dataset
        if args.clean_dataset_name != '' and args.noisy_dataset_name != '':
            # eval for model trained with mixed dataset
            print(f"using mixed dataset: {args.clean_dataset_name} and {args.noisy_dataset_name}")
            dataset = MixedSurrolDataset(args.dataset_root, args.clean_dataset_name, args.noisy_dataset_name,
                                         horizon=self.horizon,
                                         pad_before=self.pad_before, pad_after=self.pad_after, task_name=self.args.env)
        else:
            # eval for model trained with clean dataset
            dataset_name = args.clean_dataset_name if args.clean_dataset_name != '' else args.noisy_dataset_name
            print(f"using dataset: {dataset_name}")
            dataset = SurrolDataset(args.dataset_root, dataset_name, horizon=self.horizon,
                                    pad_before=self.pad_before, pad_after=self.pad_after, task_name=self.args.env)
        normalizer = dataset.get_normalizer()
        self.state_normalizer = normalizer["obs"]["state"]
        self.action_normalizer = normalizer["action"]

    def eval(self, num_episodes, render_vids=False):
        """Evaluates the model."""
        if not os.path.exists(os.path.join(args.log_eval_root, args.model_name)):
            os.makedirs(os.path.join(args.log_eval_root, args.model_name), exist_ok=True)

        rewards = []
        all_success = []
        if render_vids:
            vid_output_folder = os.path.join("videos_output", self.args.model_name)
            if not os.path.exists(vid_output_folder):
                os.makedirs(vid_output_folder)  
        for i in tqdm(range(num_episodes)):
            episode, s = self.rollout_sampler.collect_rollout(render=render_vids)
            if render_vids:

                caption = f"Env: {args.env}, Success: {episode['success'][-1]}"

                if s:
                    vid_output_path = os.path.join(vid_output_folder, 'episode'+str(i+1)+'.mp4')   
                else:
                    vid_output_path = os.path.join(vid_output_folder, '(failed)episode'+str(i+1)+'.mp4') 
                episode_to_mp4(episode, self.args.env, vid_output_path, caption=caption, fps=20)

                if self.args.isLog:
                    wandb.log({"video": wandb.Video(vid_output_path, fps=10, format="mp4")})
            r = sum(episode['reward'])
            rewards.append(r)
            all_success.append(s)
            a = int(sum(all_success))
            b = int(len(all_success))
            success_rate = int(a / b * 100)
            tqdm.write(f'episode {i}: \t reward {r}, \t success {s}, \t success rate {a}/{b}={success_rate}%')

        final_success_rate = sum(all_success)/len(all_success)
        success_rate_message = f"seed{args.seed} >>> success rate for evaluating {self.args.model_name}: {final_success_rate}\n"
        print(success_rate_message)

        # Save the success rate message to a text file
        if not args.video:
            output_file_path = os.path.join(args.log_eval_root, args.model_name, f"evaluation_results.txt")
            # output_file_path = f"/bd_byta6000i0/users/jhun/SurgicalDiffusionPolicy/seed{args.seed}_evaluation_results.txt"
            print(f"writing success rate message to {output_file_path}")
            with open(output_file_path, "a") as file:
                file.write(success_rate_message)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='args for training diffusion model')
    parser.add_argument('--env', type=str, required=True,
                        help='environment name')

    parser.add_argument('--model_root', type=str, default='trained_models',
                        help='root directory for trained models')
    parser.add_argument('--model_name', type=str, required=True,
                        help='name of the model to be evaluated')
    parser.add_argument('--dataset_root', type=str, default='datasets/demo',
                        help='root directory for datasets')
    parser.add_argument('--clean_dataset_name', type=str, default='',
                        help='name of the clean dataset, no need .npz')
    parser.add_argument('--noisy_dataset_name', type=str, default='',
                        help='name of the noisy dataset, no need .npz')
    parser.add_argument('--oaPair', type=str, default='1o1a', choices=['1o1a', '4o4a'],
                        help='use history obs, predict future action trunk')
    parser.add_argument('--log_eval_root', default='logs/eval',
                        help='root directory for evaluation logs')

    parser.add_argument('--num_episodes', type=int, default=50,
                        help='number of episodes to evaluate')
    parser.add_argument('--max_episode_length', type=int, default=100,
                        help='maximum length of an episode')

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

    parser.add_argument('--seed', type=int, default=0,
                        help='random seed for evaluation')
    parser.add_argument('--video', action='store_true',
                        help='render video during evaluation')
    args = parser.parse_args()
    print(args)

    set_seed_everywhere(args.seed)
    print(f'random seed for evaluation: {args.seed}')

    project_name = 'surrol'
    run_name = 'diffusion_eval_' + args.model_name
    if args.isLog:
        wandb.init(project=project_name, name=run_name)

    eval_sampler = EvalSampler(args)
    eval_sampler.eval(num_episodes=args.num_episodes, render_vids=args.video)
    wandb.disconnect()