test_skills.py

"""
Visualize the trained latent action space
"""

from pathlib import Path

import hydra
import numpy as np
import torch

import utils
from logger import Logger
from hydra.utils import get_original_cwd


from env_helper import make_eval_envs
torch.backends.cudnn.benchmark = True
from utils import make_agent

class Workspace:
    def __init__(self, cfg):
        self.work_dir = Path.cwd()
        print(f'workspace: {self.work_dir}')

        cfg.agent.nstep = 1
        self.cfg = cfg

        utils.set_seed_everywhere(cfg.seed)
        self.device = torch.device(cfg.device)

        # create logger
        self.logger = Logger(self.work_dir,
                             use_tb=cfg.use_tb,
                             use_wandb=cfg.use_wandb)

        self.eval_env = make_eval_envs(cfg)

        # create agent
        self.agent = make_agent(cfg.obs_type,
                                self.eval_env.observation_spec(),
                                self.eval_env.action_spec(),
                                cfg.num_seed_frames // cfg.action_repeat,
                                cfg,
                                cfg.agent)

        self.agent.init_critic = False

        self.agent = self.load_agent(self.agent)

        self.timer = utils.Timer()
        self._global_step = 0
        self._global_episode = 0

    @property
    def global_step(self):
        return self._global_step

    @property
    def global_episode(self):
        return self._global_episode

    @property
    def global_frame(self):
        return self.global_step * self.cfg.action_repeat

    def test_skill(self):
        episode, total_reward = 0, 0

        cfg = self.cfg

        print_diayn_result = cfg.print_diayn_result
        render = cfg.render

        torch.set_printoptions(sci_mode=False, precision=2)

        use_diayn = "diayn" in cfg.agent.name

        for ind in range(self.cfg.num_eval_episodes):
            step = 0
            if use_diayn:
                z = np.zeros(self.agent.skill_channel, dtype=np.int32) # Test all zero skill
                if cfg.plot_change:
                    if cfg.domain == "particle":
                        assert self.agent.skill_channel == 10
                        z = np.zeros(10, dtype=np.int32)
                    elif cfg.domain == "igibson":
                        z = np.array([1, 2, 2, 0, 2], dtype=np.int32)
                    else:
                        raise NotImplementedError(f"Unknown domain {cfg.domain} for diayn skill testing")

                print(f"testing skill {z}")
                meta = self.agent.get_meta_from_skill(z, 1)
                readable_skill = meta['skill'].reshape(self.agent.skill_channel, self.agent.skill_dim).argmax(axis=-1)
            else:
                meta = self.agent.init_meta(1)
                z = "test" # This is the file name for visualization
                meta['skill'] = np.ones_like(meta['skill']) * ind / self.cfg.num_eval_episodes
                print(f"testing skill {meta['skill']}")
            time_step = self.eval_env.reset()

            # import ipdb; ipdb.set_trace()

            while not time_step.last():
                step += 1
                if use_diayn:
                    if not cfg.plot_change:
                        meta = self.agent.update_meta(meta, step, time_step, 1, total_step=4000000)
                    else:
                        meta = meta
                    if cfg.agent.update_skill_inter_episode and step % cfg.agent.update_skill_every_step == 0:
                        if cfg.plot_change:
                            assert cfg.domain == "igibson"
                            z = np.random.randint(0, 4, size=cfg.env.igibson.gc_channel)
                            print(f"z: {z}")
                            meta = self.agent.get_meta_from_skill(z, 1)
                        readable_skill = meta['skill'].reshape(self.agent.skill_channel,
                                                               self.agent.skill_dim).argmax(axis=-1)
                        print(f"step {step} "
                              f"meta: {readable_skill}")
                else:
                    meta = self.agent.update_meta(meta, step, time_step)
                    if step % cfg.agent.update_skill_every_step == 0:
                        meta['skill'] = np.ones_like(meta['skill']) * ind / self.cfg.num_eval_episodes
                        print(f"testing skill {meta['skill']}")
                with torch.no_grad(), utils.eval_mode(self.agent):
                    action = self.agent.act(time_step.observation,
                                            meta,
                                            400000,
                                            eval_mode=True)

                action = action.flatten()

                time_step = self.eval_env.step(action)
                total_reward += time_step.reward

                if render:
                    self.eval_env.render(mode="human")

            print(f"{episode} episode done")
            episode += 1

    def load_agent(self, agent):
        snapshot_base_dir = Path(self.cfg.snapshot_base_dir)
        domain = self.cfg.domain
        snapshot_dir = snapshot_base_dir / self.cfg.obs_type / domain / self.cfg.snapshot_name

        def try_load(seed):
            actor = snapshot_dir / str(
                seed) / f'actor_{self.cfg.snapshot_ts}.pt'
            actor = get_original_cwd() / actor

            print(f"loading snapshot {actor}")
            if not actor.exists():
                return None
            with actor.open('rb') as f:
                actor = torch.load(f, map_location=self.device)
            agent.actor.load_state_dict(actor)

            if self.cfg.load_critic:
                # Now, we load critic (and diayn)
                critic = snapshot_dir / str(
                    seed) / f'critic_{self.cfg.snapshot_ts}.pt'
                critic = get_original_cwd() / critic
                with critic.open('rb') as f:
                    critic = torch.load(f, map_location=self.device)
                agent.critic.load_state_dict(critic)

            if self.cfg.load_discriminator:
                if "diayn" in self.cfg.agent.name:
                    diayn = snapshot_dir / str(
                        seed) / f'discriminator_{self.cfg.snapshot_ts}.pt'
                    diayn = get_original_cwd() / diayn
                    with diayn.open('rb') as f:
                        diayn = torch.load(f, map_location=self.device)
                    agent.diayn.load_state_dict(diayn)

                    if self.cfg.agent.anti:
                        anti_diayn = snapshot_dir / str(
                            seed) / f'anti_{self.cfg.snapshot_ts}.pt'
                        anti_diayn = get_original_cwd() / anti_diayn
                        with anti_diayn.open('rb') as f:
                            anti_diayn = torch.load(f, map_location=self.device)
                        agent.anti_diayn.load_state_dict(anti_diayn)
            return agent

        # try to load current seed
        payload = try_load(self.cfg.seed)
        if payload is not None:
            return payload
        # otherwise try random seed
        while True:
            seed = np.random.randint(1, 11)
            payload = try_load(seed)
            if payload is not None:
                return payload
        return None

@hydra.main(config_path='.', config_name='test_skills')
def main(cfg):
    root_dir = Path.cwd()
    workspace = Workspace(cfg)
    snapshot = root_dir / 'snapshot.pt'
    if snapshot.exists():
        print(f'resuming: {snapshot}')
        workspace.load_snapshot()
    workspace.test_skill()


if __name__ == '__main__':
    main()