eval.py

import os
#os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import pathlib
import numpy as np
import click
import json
import jittor
import taichi as ti
import random

from shoot import MPMSolver
from env import JellyEnv
from Autoencoder import AutoEncoder
from rlkit.jittor.sac.policies import TanhGaussianPolicy
from rlkit.jittor.networks import FlattenMlp, MlpEncoder
from rlkit.jittor.sac.sac import PEARLSoftActorCritic
from rlkit.jittor.sac.agent import PEARLAgent
from rlkit.launchers.launcher_util import setup_logger
import rlkit.jittor.jittor_util as ptu
from configs.default import default_config
from rlkit.jittor.sac.policies import MakeDeterministic
import torch
from sklearn.model_selection import train_test_split
import torch
from torch import nn
import torch.nn.functional as F
import torch.utils.data as Data 
import torch.optim as optim
from torch.utils.data import Dataset

gui=ti.GUI("MPM SCOOP THE JELLY", res=(512,512), background_color=0x112F41,show_gui=True)

def setup_seed(seed=1024):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    jittor.set_global_seed(seed) 

setup_seed() 

if jittor.has_cuda:
    jittor.flags.use_cuda = 1

def experiment(variant):
    env = JellyEnv(1, 'model/autoencoder_full.pth', gui)
    deterministic = True
    tasks = env.get_all_task_idx()
    obs_dim = int(np.prod(env.observation_space.shape))
    action_dim = int(np.prod(env.action_space.shape))
    reward_dim = 1

    # instantiate networks
    latent_dim = variant['latent_size']
    context_encoder_input_dim = 2 * obs_dim + action_dim + reward_dim
    context_encoder_output_dim = latent_dim * 2
    net_size = variant['net_size']
    encoder_model = MlpEncoder

    context_encoder = encoder_model(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=context_encoder_input_dim,
        output_size=context_encoder_output_dim,
    )
    qf1 = FlattenMlp(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=obs_dim + action_dim + latent_dim,
        output_size=1,
    )
    qf2 = FlattenMlp(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=obs_dim + action_dim + latent_dim,
        output_size=1,
    )
    vf = FlattenMlp(
        hidden_sizes=[net_size, net_size, net_size],
        input_size=obs_dim + latent_dim,
        output_size=1,
    )
    policy = TanhGaussianPolicy(
        hidden_sizes=[net_size, net_size, net_size],
        obs_dim=obs_dim + latent_dim,
        latent_dim=latent_dim,
        action_dim=action_dim,
    )
    agent = PEARLAgent(
        latent_dim,
        context_encoder,
        policy,
        **variant['algo_params']
    )

    algorithm = PEARLSoftActorCritic(
        env=env,
        train_tasks=list(tasks[:variant['n_train_tasks']]),
        eval_tasks=list(tasks[-variant['n_eval_tasks']:]),
        nets=[agent, qf1, qf2, vf],
        latent_dim=latent_dim,
        **variant['algo_params']
    )


    print('initialize state')
    algorithm.eval(10)
    qf1.load('model/qf1.pkl')
    qf2.load('model/qf2.pkl')
    vf.load('model/vf.pkl')
    policy.load('model/policy.pkl')
    context_encoder.load('model/context_encoder.pkl')
    print('------------------------------------------------')
    print('eval state')
    algorithm.eval(10)
    
def deep_update_dict(fr, to):
    ''' update dict of dicts with new values '''
    # assume dicts have same keys
    for k, v in fr.items():
        if type(v) is dict:
            deep_update_dict(v, to[k])
        else:
            to[k] = v
    return to

@click.command()
@click.argument('config', default=None)
def main(config):
    variant = default_config
    if config:
        with open(os.path.join(config)) as f:
            exp_params = json.load(f)
        variant = deep_update_dict(exp_params, variant)

    experiment(variant)

if __name__ == "__main__":
    main()