|
|
import numpy as np
|
|
|
import random
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import os
|
|
|
import inspect
|
|
|
import pickle
|
|
|
import gdown
|
|
|
from network import Actor
|
|
|
|
|
|
|
|
|
def weight_init(m):
|
|
|
"""Custom weight init for Conv2D and Linear layers.
|
|
|
Reference: https://github.com/MishaLaskin/rad/blob/master/curl_sac.py"""
|
|
|
|
|
|
if isinstance(m, nn.Linear):
|
|
|
nn.init.orthogonal_(m.weight.data)
|
|
|
m.bias.data.fill_(0.0)
|
|
|
elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
|
|
|
|
|
|
assert m.weight.size(2) == m.weight.size(3)
|
|
|
m.weight.data.fill_(0.0)
|
|
|
m.bias.data.fill_(0.0)
|
|
|
mid = m.weight.size(2) // 2
|
|
|
gain = nn.init.calculate_gain('relu')
|
|
|
nn.init.orthogonal_(m.weight.data[:, :, mid, mid], gain)
|
|
|
|
|
|
|
|
|
def set_seed(random_seed):
|
|
|
if random_seed <= 0:
|
|
|
random_seed = np.random.randint(1, 9999)
|
|
|
else:
|
|
|
random_seed = random_seed
|
|
|
|
|
|
torch.manual_seed(random_seed)
|
|
|
np.random.seed(random_seed)
|
|
|
random.seed(random_seed)
|
|
|
|
|
|
return random_seed
|
|
|
|
|
|
|
|
|
def make_env(env_name, seed):
|
|
|
import gymnasium as gym
|
|
|
|
|
|
env = gym.make(env_name)
|
|
|
env.action_space.seed(seed)
|
|
|
state_dim = env.observation_space.shape[0]
|
|
|
action_dim = env.action_space.shape[0]
|
|
|
action_bound = [env.action_space.low[0], env.action_space.high[0]]
|
|
|
|
|
|
env_info = {'name': env_name, 'state_dim': state_dim, 'action_dim': action_dim, 'action_bound': action_bound, 'seed': seed}
|
|
|
|
|
|
return env, env_info
|
|
|
|
|
|
|
|
|
def get_learning_info(args, seed):
|
|
|
env, env_info = make_env(args.env_name, seed)
|
|
|
device = 'cuda'
|
|
|
|
|
|
alpha_dict = {'HalfCheetah-v3': args.alpha_threshold, 'Walker2d-v3': args.alpha_threshold,
|
|
|
'Ant-v3': args.alpha_threshold, 'Hopper-v3': args.alpha_threshold}
|
|
|
|
|
|
thresholds = {"ALPHA_THRESHOLD": alpha_dict[args.env_name], "THETA_THRESHOLD": args.theta_threshold}
|
|
|
max_action = 1
|
|
|
|
|
|
t_p = Actor(env_info['state_dim'], env_info['action_dim'], (400, 300), 1)
|
|
|
num_teacher_param = sum(p2.numel() for p2 in t_p.parameters())
|
|
|
|
|
|
kwargs = {
|
|
|
"env": env,
|
|
|
"args": args,
|
|
|
"env_info": env_info,
|
|
|
"thresholds": thresholds,
|
|
|
"discount": args.discount,
|
|
|
"datasize": args.datasize,
|
|
|
"tau": args.tau,
|
|
|
"device": device,
|
|
|
"num_teacher_param": num_teacher_param,
|
|
|
"noise_clip": args.noise_clip * max_action,
|
|
|
"policy_freq": args.policy_freq,
|
|
|
"h": args.h,
|
|
|
}
|
|
|
return kwargs
|
|
|
|
|
|
|
|
|
def get_compression_ratio(num_teacher_param, agent):
|
|
|
kep_w = 0
|
|
|
for c in agent.actor.children():
|
|
|
kep_w += c.get_num_remained_weights()
|
|
|
|
|
|
|
|
|
return kep_w / num_teacher_param
|
|
|
|
|
|
|
|
|
def load_buffer(env_name, level, datasize):
|
|
|
current_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
|
|
|
file_path = os.path.join(current_dir, "teacher_buffer", "[" + level + "_buffer]_" + env_name + ".pickle")
|
|
|
try:
|
|
|
with open(file_path, "rb") as fr:
|
|
|
buffer = pickle.load(fr)
|
|
|
buffer.size = datasize
|
|
|
except FileNotFoundError:
|
|
|
|
|
|
if level == 'expert':
|
|
|
print("Downloading the teacher buffer...")
|
|
|
if env_name == "Ant-v3":
|
|
|
file_id = "10VBf3bM38bNw9WsniQvirpNjRFWp8HZO"
|
|
|
elif env_name == "Walker2d-v3":
|
|
|
file_id = "1ungLoqNKS4NIldZ9H2mswwGh-3Ipgy0D"
|
|
|
elif env_name == "HalfCheetah-v3":
|
|
|
file_id = "1wO0HwDi1GNf9d9SrDJrf9x8XMZDOTkzl"
|
|
|
elif env_name == "Hopper-v3":
|
|
|
file_id ="10pqCliJSM_Iyb05dxHZfYs9VlmCmPryE"
|
|
|
else:
|
|
|
raise ValueError("Invalid Environment Name")
|
|
|
|
|
|
url = f"https://drive.google.com/uc?id={file_id}"
|
|
|
gdown.download(url, file_path, quiet=False)
|
|
|
print("Download Complete!")
|
|
|
elif level == 'medium':
|
|
|
if env_name == "Ant-v3":
|
|
|
file_id = "1-SKleNu6l-tY2awkx3tgVDUKbjkOaj_D"
|
|
|
elif env_name == "Walker2d-v3":
|
|
|
file_id = "1x6nkBBSWMRb3bENxUzcntHT1WlSNJmoh"
|
|
|
elif env_name == "HalfCheetah-v3":
|
|
|
file_id = "1OHkB6yVK3QcqbuJH0B_iNW_2cBnv96mR"
|
|
|
elif env_name == "Hopper-v3":
|
|
|
file_id ="1uqH2pgKKrhadsCXCwQWrvDvZ4ZyYFkM-"
|
|
|
else:
|
|
|
raise ValueError("Invalid Environment Name")
|
|
|
|
|
|
url = f"https://drive.google.com/uc?id={file_id}"
|
|
|
gdown.download(url, file_path, quiet=False)
|
|
|
|
|
|
else:
|
|
|
raise ValueError("Invalid Level. Choose from ['expert', 'medium']")
|
|
|
|
|
|
with open(file_path, "rb") as fr:
|
|
|
buffer = pickle.load(fr)
|
|
|
buffer.size = datasize
|
|
|
|
|
|
return buffer
|
|
|
|
