def __init__(self,

cfg=drnd_config) -> None:

make_dir("weights")

self.device = cfg.device

self.batch_size = cfg.batch_size

self.cfg = cfg

# self.eval_env = gym.make(cfg.dataset_name)

# self.eval_env.seed(cfg.eval_seed)

# d4rl_dataset = d4rl.qlearning_dataset(self.eval_env)

# self.state_dim = self.eval_env.observation_space.shape[0]

# self.action_dim = self.eval_env.action_space.shape[0]

self.state_dim = cfg.state_dim

self.action_dim = cfg.action_dim

self.buffer = ReplayBuffer(self.state_dim, self.action_dim)

# self.buffer.from_d4rl(d4rl_dataset)

self.buffer.from_json(cfg.dataset_name)

seed_everything(cfg.train_seed)

def train_drnd(self) -> DRND:

(state_mean, state_std), (action_mean, action_std) = self.buffer.get_moments()

self.state_mean = state_mean.to(self.device)

self.state_std = state_std.to(self.device)

self.action_mean = action_mean.to(self.device)

self.action_std = action_std.to(self.device)

drnd = DRND(self.state_dim,

self.action_dim,

self.cfg.drnd_embedding_dim,

self.cfg.drnd_num_targets,

self.cfg.drnd_alpha,

self.state_mean,

self.state_std,

self.action_mean,

self.action_std,

hidden_dim=self.cfg.drnd_hidden_dim).to(self.device)

drnd_optim = Adam(drnd.predictor.parameters(), lr=self.cfg.drnd_learning_rate)

for epoch in trange(self.cfg.drnd_num_epochs, desc="DRND Epochs"):

for _ in trange(self.cfg.num_updates_on_epoch, desc="DRND Iterations"):

states, actions, _, _, _, = self.buffer.sample(self.batch_size)

states, actions = [x.to(self.device) for x in (states, actions)]

loss, update_info = drnd.update_drnd(states, actions)

drnd_optim.zero_grad()

loss.backward()

drnd_optim.step()

wandb.log(update_info)

return drnd

def train(self):

'''

run_name = f"sac_drnd_" + str(self.cfg.train_seed)

print(f"Training starts on {self.cfg.device} 🚀")

with wandb.init(project=self.cfg.project, group=self.cfg.group, name=run_name, job_type="offline_training"):

wandb.config.update({k: v for k, v in self.cfg.__dict__.items() if not k.startswith("__")})

drnd = self.train_drnd()

drnd.eval()

actor = Actor(self.state_dim, self.action_dim, self.cfg.hidden_dim)

actor_optim = Adam(actor.parameters(), lr=self.cfg.actor_lr)

critic = EnsembledCritic(self.state_dim, self.action_dim, self.cfg.hidden_dim, layer_norm=self.cfg.critic_layernorm)

critic_optim = Adam(critic.parameters(), lr=self.cfg.critic_lr)

self.sac_drnd = SAC_DRND(actor,

actor_optim,

critic,

critic_optim,

drnd,

self.cfg.actor_lambda,

self.cfg.critic_lambda,

self.cfg.beta_lr,

self.cfg.gamma,

self.cfg.tau,

self.device)

for t in tqdm(range(self.cfg.max_timesteps), desc="SAC DRND steps"):

batch = self.buffer.sample(self.batch_size)

states, actions, rewards, next_states, dones = [x.to(self.device) for x in batch]

logging_dict = self.sac_drnd.train_offline_step(states,

actions,

rewards,

next_states,

dones)

wandb.log(logging_dict, step=self.sac_drnd.total_iterations)

# for epoch in trange(self.cfg.num_epochs, desc="Offline SAC Epochs"):

# update_info_total = {

# "sac_offline/actor_loss": 0,

# "sac_offline/actor_batch_entropy": 0,

# "sac_offline/drnd_policy": 0,

# "sac_offline/drnd_random": 0,

# "sac_offline/critic_loss": 0,

# "sac_offline/q_mean": 0

# }

# for _ in range(self.cfg.num_updates_on_epoch):

# state, action, reward, next_state, done = self.buffer.sample(self.batch_size)

# update_info = self.sac_drnd.train_offline_step(state,

# action,

# reward,

# next_state,

# done)

# for k, v in update_info.items():

# update_info_total[k] += v

# for k, v in update_info_total.items():

# update_info_total[k] /= self.cfg.num_updates_on_epoch

# wandb.log(update_info)

# if epoch % self.cfg.eval_period == 0 or epoch == self.cfg.num_epochs - 1:

# eval_returns = self.eval_actor()

# normalized_score = self.eval_env.get_normalized_score(eval_returns) * 100.0

# wandb.log({

# "eval/return_mean": np.mean(eval_returns),

# "eval/return_std": np.std(eval_returns),

# "eval/normalized_score_mean": np.mean(normalized_score),

# "eval/normalized_score_std": np.std(normalized_score)

# })

wandb.finish()

# @torch.no_grad()

# def eval_actor(self) -> np.ndarray:

# self.eval_env.seed(self.cfg.eval_seed)

# self.sac_drnd.actor.eval()

# episode_rewards = []

# for _ in range(self.cfg.eval_episodes):

# state, done = self.eval_env.reset(), False

# episode_reward = 0.0

# while not done:

# action = self.sac_drnd.actor.act(state, self.device)

# state, reward, done, _ = self.eval_env.step(action)

# episode_reward += reward

# episode_rewards.append(episode_reward)

# self.sac_drnd.actor.train()

# return np.array(episode_rewards)

def save(self):

state_dict = self.sac_drnd.state_dict()

torch.save(state_dict, self.cfg.checkpoint_path)

# def load(self, map_location: str = "cpu"):

# state_dict = torch.load(self.cfg.checkpoint_path, map_location=map_location)

# actor = Actor(self.state_dim, self.action_dim, self.cfg.hidden_dim)

# actor_optim = Adam(actor.parameters(), lr=self.cfg.actor_lr)

# critic = EnsembledCritic(self.state_dim, self.action_dim, self.cfg.hidden_dim, layer_norm=self.cfg.critic_layernorm)

# critic_optim = Adam(critic.parameters(), lr=self.cfg.critic_lr)

# rnd = DRND(self.state_dim,

# self.action_dim,

# self.cfg.rnd_embedding_dim,

# self.state_mean,

# self.state_std,

# self.action_mean,

# self.action_std,

# hidden_dim=self.cfg.rnd_hidden_dim)

# self.sac_drnd = SAC_DRND(actor,

# actor_optim,

# critic,

# critic_optim,

# rnd,

# self.cfg.actor_alpha,

# self.cfg.critic_alpha,

# self.cfg.beta_lr,

# self.cfg.gamma,

# self.cfg.tau,

# self.device)