| import copy |
| import torch |
| import torch.nn.functional as F |
| from SGVLB import SGVLB |
| from network import Net, Critic |
|
|
|
|
| class BPDAgent(object): |
| def __init__( |
| self, |
| env, |
| args, |
| env_info, |
| thresholds, |
| datasize, |
| device, |
| discount, |
| tau, |
| noise_clip, |
| policy_freq, |
| h, |
| num_teacher_param, |
| ): |
| self.args = args |
| self.env = env |
| self.env_info = env_info |
|
|
| self.actor = Net(env_info['state_dim'], env_info['action_dim'], env_info['action_bound'], |
| args.student_hidden_dims, thresholds['ALPHA_THRESHOLD'], thresholds['THETA_THRESHOLD'], |
| device=device).to(device) |
| self.actor_target = copy.deepcopy(self.actor) |
| self.sgvlb = SGVLB(self.actor, datasize, loss_type='l2', device=device) |
| self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=3e-4) |
|
|
| self.critic = Critic(env_info['state_dim'], env_info['action_dim']).to(device) |
| self.critic_target = copy.deepcopy(self.critic) |
| self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=3e-4) |
|
|
| self.discount = discount |
| self.tau = tau |
| self.noise_clip = noise_clip |
| self.policy_freq = policy_freq |
| self.datasize = datasize |
| self.h = h |
|
|
| self.total_it = 0 |
| self.kl_weight = 0 |
|
|
| def set_kl_weight(self, kl_weight): |
| self.kl_weight = kl_weight |
| return |
|
|
| def test(self): |
| self.actor.eval() |
| with torch.no_grad(): |
| return_list = [] |
| for epi_cnt in range(1, self.args.num_test_epi): |
| episode_return = 0 |
| done = False |
| state, _ = self.env.reset() |
| while not done: |
| action = self.actor(state) |
| action = action.cpu().numpy()[0] |
| next_state, reward, terminated, truncated, _ = self.env.step(action) |
| done = terminated or truncated |
| episode_return += reward |
| state = next_state |
| return_list.append(episode_return) |
|
|
| avg_return = sum(return_list) / len(return_list) |
| max_return = max(return_list) |
| min_return = min(return_list) |
|
|
| return avg_return, max_return, min_return |
|
|
| def train(self, transition): |
| self.actor.train() |
|
|
| self.total_it += 1 |
|
|
| states, actions, rewards, next_states, dones = transition |
|
|
| with torch.no_grad(): |
| next_actions = ( |
| self.actor_target(next_states) |
| ).clamp(self.env_info['action_bound'][0], self.env_info['action_bound'][1]) |
|
|
| target_Q1, target_Q2 = self.critic_target(next_states, next_actions) |
| target_Q = torch.min(target_Q1, target_Q2) |
| target_Q = rewards + (1 - dones) * self.discount * target_Q |
|
|
| current_Q1, current_Q2 = self.critic(states, actions) |
| critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q) |
|
|
| self.critic_optimizer.zero_grad() |
| critic_loss.backward() |
| self.critic_optimizer.step() |
|
|
| if self.total_it % self.policy_freq == 0: |
| pi = self.actor(states) |
| Q = self.critic.Q1(states, pi) |
| lmbda = (self.h * self.datasize) / Q.abs().mean().detach() |
|
|
| actor_loss = -lmbda * Q.mean() + self.sgvlb(pi, actions, self.kl_weight) |
|
|
| |
| self.actor_optimizer.zero_grad() |
| actor_loss.backward() |
| self.actor_optimizer.step() |
|
|
| |
| for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()): |
| target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) |
|
|
| for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()): |
| target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) |
|
|
| def __del__(self): |
| del self.actor |
| del self.actor_target |
| del self.critic |
| del self.critic_target |
| return |
|
|
|
|
