mirror of
https://github.com/wassname/pytorch-soft-actor-critic.git
synced 2026-07-16 11:20:55 +08:00
108 lines
3.3 KiB
Python
108 lines
3.3 KiB
Python
import sys
|
|
import os
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.distributions import Normal
|
|
from utils import create_log_gaussian, logsumexp
|
|
|
|
LOG_SIG_MAX = 2
|
|
LOG_SIG_MIN = -20
|
|
epsilon=1e-6
|
|
|
|
# Initialize Policy weights
|
|
def weights_init_policy(m):
|
|
classname = m.__class__.__name__
|
|
if classname.find('Linear') != -1:
|
|
torch.nn.init.normal_(m.weight, mean=0, std=0.1)
|
|
|
|
# Initialize QNetwork and Value Network weights
|
|
def weights_init_vf(m):
|
|
classname = m.__class__.__name__
|
|
if classname.find('Linear') != -1:
|
|
torch.nn.init.xavier_normal_(m.weight)
|
|
|
|
|
|
class ValueNetwork(nn.Module):
|
|
def __init__(self, state_dim, hidden_dim):
|
|
super(ValueNetwork, self).__init__()
|
|
|
|
self.linear1 = nn.Linear(state_dim, hidden_dim)
|
|
self.linear2 = nn.Linear(hidden_dim, hidden_dim)
|
|
self.linear3 = nn.Linear(hidden_dim, 1)
|
|
|
|
self.apply(weights_init_vf)
|
|
|
|
def forward(self, state):
|
|
x = F.relu(self.linear1(state))
|
|
x = F.relu(self.linear2(x))
|
|
x = self.linear3(x)
|
|
return x
|
|
|
|
|
|
class QNetwork(nn.Module):
|
|
def __init__(self, num_inputs, num_actions, hidden_size):
|
|
super(QNetwork, self).__init__()
|
|
|
|
# Q1 architecture
|
|
self.linear1 = nn.Linear(num_inputs + num_actions, hidden_size)
|
|
self.linear2 = nn.Linear(hidden_size, hidden_size)
|
|
self.linear3 = nn.Linear(hidden_size, 1)
|
|
|
|
# Q2 architecture
|
|
self.linear4 = nn.Linear(num_inputs + num_actions, hidden_size)
|
|
self.linear5 = nn.Linear(hidden_size, hidden_size)
|
|
self.linear6 = nn.Linear(hidden_size, 1)
|
|
|
|
self.apply(weights_init_vf)
|
|
|
|
def forward(self, state, action):
|
|
x1 = torch.cat([state, action], 1)
|
|
x1 = F.relu(self.linear1(x1))
|
|
x1 = F.relu(self.linear2(x1))
|
|
x1 = self.linear3(x1)
|
|
|
|
x2 = torch.cat([state, action], 1)
|
|
x2 = F.relu(self.linear4(x2))
|
|
x2 = F.relu(self.linear5(x2))
|
|
x2 = self.linear6(x2)
|
|
|
|
return x1, x2
|
|
|
|
|
|
class GaussianPolicy(nn.Module):
|
|
def __init__(self, num_inputs, num_actions, hidden_size):
|
|
super(GaussianPolicy, self).__init__()
|
|
|
|
self.linear1 = nn.Linear(num_inputs, hidden_size)
|
|
self.linear2 = nn.Linear(hidden_size, hidden_size)
|
|
|
|
self.mean_linear = nn.Linear(hidden_size, num_actions)
|
|
|
|
self.log_std_linear = nn.Linear(hidden_size, num_actions)
|
|
|
|
self.apply(weights_init_policy)
|
|
|
|
def forward(self, state):
|
|
x = F.relu(self.linear1(state))
|
|
x = F.relu(self.linear2(x))
|
|
mean = self.mean_linear(x)
|
|
log_std = self.log_std_linear(x)
|
|
log_std = torch.clamp(log_std, min=LOG_SIG_MIN, max=LOG_SIG_MAX)
|
|
return mean, log_std
|
|
|
|
def evaluate(self, state, reparam=False):
|
|
mean, log_std = self.forward(state)
|
|
std = log_std.exp()
|
|
normal = Normal(mean, std)
|
|
if reparam == True:
|
|
x_t = normal.rsample() # reparameterization trick (mean + std * N(0,1))
|
|
else:
|
|
x_t = normal.sample() # log-derivative trick (N(mean, std))
|
|
action = torch.tanh(x_t)
|
|
log_prob = normal.log_prob(x_t)
|
|
# Enforcing Action Bound
|
|
log_prob -= torch.log(1 - action.pow(2) + epsilon)
|
|
log_prob = log_prob.sum(-1, keepdim=True)
|
|
return action, log_prob, x_t, mean, log_std
|