add

Author: xuhaoran4

BCQ_L.py

@@ -0,0 +1,203 @@
+import copy
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Actor(nn.Module):
+    def __init__(self, state_dim, action_dim, max_action, hidden_unit=256, phi=0.05):
+        super(Actor, self).__init__()
+        self.l1 = nn.Linear(state_dim + action_dim, hidden_unit)
+        self.l2 = nn.Linear(hidden_unit, hidden_unit)
+        self.l3 = nn.Linear(hidden_unit, action_dim)
+
+        self.max_action = max_action
+        self.phi = phi
+
+    def forward(self, state, action):
+        a = F.relu(self.l1(torch.cat([state, action], 1)))
+        a = F.relu(self.l2(a))
+        a = self.phi * self.max_action * torch.tanh(self.l3(a))
+        return (a + action).clamp(-self.max_action, self.max_action)
+
+
+class Critic(nn.Module):
+    def __init__(self, state_dim, action_dim, hidden_unit=256):
+        super(Critic, self).__init__()
+        self.l1 = nn.Linear(state_dim + action_dim, hidden_unit)
+        self.l2 = nn.Linear(hidden_unit, hidden_unit)
+        self.l3 = nn.Linear(hidden_unit, 1)
+
+        self.l4 = nn.Linear(state_dim + action_dim, hidden_unit)
+        self.l5 = nn.Linear(hidden_unit, hidden_unit)
+        self.l6 = nn.Linear(hidden_unit, 1)
+
+    def forward(self, state, action):
+        q1 = F.relu(self.l1(torch.cat([state, action], 1)))
+        q1 = F.relu(self.l2(q1))
+        q1 = self.l3(q1)
+
+        q2 = F.relu(self.l4(torch.cat([state, action], 1)))
+        q2 = F.relu(self.l5(q2))
+        q2 = self.l6(q2)
+        return q1, q2
+
+    def q1(self, state, action):
+        q1 = F.relu(self.l1(torch.cat([state, action], 1)))
+        q1 = F.relu(self.l2(q1))
+        q1 = self.l3(q1)
+        return q1
+
+
+# Vanilla Variational Auto-Encoder 
+class VAE(nn.Module):
+    def __init__(self, state_dim, action_dim, latent_dim, max_action, hidden_unit=256):
+        super(VAE, self).__init__()
+        self.e1 = nn.Linear(state_dim + action_dim, hidden_unit)
+        self.e2 = nn.Linear(hidden_unit, hidden_unit)
+
+        self.mean = nn.Linear(hidden_unit, latent_dim)
+        self.log_std = nn.Linear(hidden_unit, latent_dim)
+
+        self.d1 = nn.Linear(state_dim + latent_dim, hidden_unit)
+        self.d2 = nn.Linear(hidden_unit, hidden_unit)
+        self.d3 = nn.Linear(hidden_unit, action_dim)
+
+        self.max_action = max_action
+        self.latent_dim = latent_dim
+
+    def forward(self, state, action):
+        z = F.relu(self.e1(torch.cat([state, action], 1)))
+        z = F.relu(self.e2(z))
+
+        mean = self.mean(z)
+        # Clamped for numerical stability
+        log_std = self.log_std(z).clamp(-4, 15)
+        std = torch.exp(log_std)
+        z = mean + std * torch.randn_like(std)
+
+        u = self.decode(state, z)
+
+        return u, mean, std
+
+    def decode(self, state, z=None):
+        # When sampling from the VAE, the latent vector is clipped to [-0.5, 0.5]
+        if z is None:
+            z = torch.randn((state.shape[0], self.latent_dim)).to(device).clamp(-0.5, 0.5)
+
+        a = F.relu(self.d1(torch.cat([state, z], 1)))
+        a = F.relu(self.d2(a))
+        return self.max_action * torch.tanh(self.d3(a))
+
+
+class BCQ(object):
+    def __init__(self, state_dim, action_dim, max_action, discount=0.99, tau=0.005, lmbda=0.75, phi=0.05):
+        latent_dim = action_dim * 2
+
+        self.actor = Actor(state_dim, action_dim, max_action, phi=phi).to(device)
+        self.actor_target = copy.deepcopy(self.actor)
+        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=1e-3)
+
+        self.critic = Critic(state_dim, action_dim).to(device)
+        self.critic_target = copy.deepcopy(self.critic)
+        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=1e-3)
+
+        self.cost_critic = Critic(state_dim, action_dim).to(device)
+        self.cost_critic_target = copy.deepcopy(self.cost_critic)
+        self.cost_critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=1e-3)
+
+        self.vae = VAE(state_dim, action_dim, latent_dim, max_action).to(device)
+        self.vae_optimizer = torch.optim.Adam(self.vae.parameters())
+
+        self.max_action = max_action
+        self.action_dim = action_dim
+        self.discount = discount
+        self.tau = tau
+        self.lmbda = lmbda
+
+    def select_action(self, state):
+        with torch.no_grad():
+            state = torch.FloatTensor(state.reshape(1, -1)).repeat(100, 1).to(device)
+            action = self.actor(state, self.vae.decode(state))
+            q1 = self.critic.q1(state, action)
+            ind = q1.argmax(0)
+        return action[ind].cpu().data.numpy().flatten()
+
+    def train(self, replay_buffer, batch_size=100):
+        # Sample replay buffer / batch
+        state, action, next_state, reward, not_done = replay_buffer.sample(batch_size)
+
+        # Variational Auto-Encoder Training
+        recon, mean, std = self.vae(state, action)
+        recon_loss = F.mse_loss(recon, action)
+        KL_loss	= -0.5 * (1 + torch.log(std.pow(2)) - mean.pow(2) - std.pow(2)).mean()
+        vae_loss = recon_loss + 0.5 * KL_loss
+
+        self.vae_optimizer.zero_grad()
+        vae_loss.backward()
+        self.vae_optimizer.step()
+
+        # Reward Critic Training
+        with torch.no_grad():
+            # Duplicate next state 10 times
+            next_state = torch.repeat_interleave(next_state, 10, 0)
+
+            # Compute value of perturbed actions sampled from the VAE
+            target_Q1, target_Q2 = self.critic_target(next_state, self.actor_target(next_state, self.vae.decode(next_state)))
+
+            # Soft Clipped Double Q-learning
+            target_Q = self.lmbda * torch.min(target_Q1, target_Q2) + (1. - self.lmbda) * torch.max(target_Q1, target_Q2)
+            # Take max over each action sampled from the VAE
+            target_Q = target_Q.reshape(batch_size, -1).max(1)[0].reshape(-1, 1)
+
+            target_Q = reward + not_done * self.discount * target_Q
+
+        current_Q1, current_Q2 = self.critic(state, action)
+        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()
+
+        # Cost Critic Training
+        with torch.no_grad():
+            # Duplicate next state 10 times
+            next_state = torch.repeat_interleave(next_state, 10, 0)
+
+            # Compute value of perturbed actions sampled from the VAE
+            target_Q1, target_Q2 = self.critic_target(next_state, self.actor_target(next_state, self.vae.decode(next_state)))
+
+            # Soft Clipped Double Q-learning
+            target_Q = self.lmbda * torch.min(target_Q1, target_Q2) + (1. - self.lmbda) * torch.max(target_Q1, target_Q2)
+            # Take max over each action sampled from the VAE
+            target_Q = target_Q.reshape(batch_size, -1).max(1)[0].reshape(-1, 1)
+
+            target_Q = reward + not_done * self.discount * target_Q
+
+        current_Q1, current_Q2 = self.critic(state, action)
+        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()
+
+        # Pertubation Model / Action Training
+        sampled_actions = self.vae.decode(state)
+        perturbed_actions = self.actor(state, sampled_actions)
+
+        # Update through DPG
+        actor_loss = -self.critic.q1(state, perturbed_actions).mean()
+
+        self.actor_optimizer.zero_grad()
+        actor_loss.backward()
+        self.actor_optimizer.step()
+
+        # Update Target Networks
+        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)