File size: 3,565 Bytes
f761808 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from agent.helpers import SinusoidalPosEmb, init_weights
class Critic(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim=256):
super(Critic, self).__init__()
self.q1_model = nn.Sequential(nn.Linear(state_dim + action_dim, hidden_dim),
nn.Mish(),
nn.Linear(hidden_dim, hidden_dim),
nn.Mish(),
nn.Linear(hidden_dim, hidden_dim),
nn.Mish(),
nn.Linear(hidden_dim, 1))
self.q2_model = nn.Sequential(nn.Linear(state_dim + action_dim, hidden_dim),
nn.Mish(),
nn.Linear(hidden_dim, hidden_dim),
nn.Mish(),
nn.Linear(hidden_dim, hidden_dim),
nn.Mish(),
nn.Linear(hidden_dim, 1))
self.apply(init_weights)
def forward(self, state, action):
x = torch.cat([state, action], dim=-1)
return self.q1_model(x), self.q2_model(x)
def q1(self, state, action):
x = torch.cat([state, action], dim=-1)
return self.q1_model(x)
def q_min(self, state, action):
q1, q2 = self.forward(state, action)
return torch.min(q1, q2)
class Model(nn.Module):
def __init__(self, state_dim, action_dim, hidden_size=256, time_dim=32):
super(Model, self).__init__()
self.time_mlp = nn.Sequential(
SinusoidalPosEmb(time_dim),
nn.Linear(time_dim, hidden_size),
nn.Mish(),
nn.Linear(hidden_size, time_dim),
)
input_dim = state_dim + action_dim + time_dim
self.layer = nn.Sequential(nn.Linear(input_dim, hidden_size),
nn.Mish(),
nn.Linear(hidden_size, hidden_size),
nn.Mish(),
nn.Linear(hidden_size, hidden_size),
nn.Mish(),
nn.Linear(hidden_size, action_dim))
self.apply(init_weights)
def forward(self, x, time, state):
t = self.time_mlp(time)
out = torch.cat([x, t, state], dim=-1)
out = self.layer(out)
return out
class MLP(nn.Module):
def __init__(self, state_dim, action_dim, hidden_size=256):
super(MLP, self).__init__()
input_dim = state_dim
self.mid_layer = nn.Sequential(nn.Linear(input_dim, hidden_size),
nn.Mish(),
nn.Linear(hidden_size, hidden_size),
nn.Mish(),
nn.Linear(hidden_size, hidden_size),
nn.Mish())
self.final_layer = nn.Linear(hidden_size, action_dim)
self.apply(init_weights)
def forward(self, state, eval=False):
out = self.mid_layer(state)
out = self.final_layer(out)
if not eval:
out += torch.randn_like(out) * 0.1
return out
def loss(self, action, state):
return F.mse_loss(self.forward(state), action, reduction='mean')
|