wehub-resource-sync
51 lines
1.3 KiB
Python
51 lines
1.3 KiB
Python
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from agents.helpers import SinusoidalPosEmb
|
|
|
|
|
|
class MLP(nn.Module):
|
|
"""
|
|
MLP Model
|
|
"""
|
|
def __init__(self,
|
|
state_dim,
|
|
action_dim,
|
|
device,
|
|
t_dim=16):
|
|
|
|
super(MLP, self).__init__()
|
|
self.device = device
|
|
|
|
self.time_mlp = nn.Sequential(
|
|
SinusoidalPosEmb(t_dim),
|
|
nn.Linear(t_dim, t_dim * 2),
|
|
nn.Mish(),
|
|
nn.Linear(t_dim * 2, t_dim),
|
|
)
|
|
|
|
input_dim = state_dim + action_dim + t_dim
|
|
self.mid_layer = nn.Sequential(nn.Linear(input_dim, 256),
|
|
nn.Mish(),
|
|
nn.Linear(256, 256),
|
|
nn.Mish(),
|
|
nn.Linear(256, 256),
|
|
nn.Mish())
|
|
|
|
self.final_layer = nn.Linear(256, action_dim)
|
|
|
|
def forward(self, x, time, state):
|
|
|
|
t = self.time_mlp(time)
|
|
x = torch.cat([x, t, state], dim=1)
|
|
x = self.mid_layer(x)
|
|
|
|
return self.final_layer(x)
|
|
|
|
|