in monobeast/minigrid/monobeast_amigo.py [0:0]
def __init__(self, observation_shape, num_actions, state_embedding_dim=256, num_input_frames=1, use_lstm=False, num_lstm_layers=1):
super(MinigridNet, self).__init__()
self.observation_shape = observation_shape
self.num_actions = num_actions
self.state_embedding_dim = state_embedding_dim
self.use_lstm = use_lstm
self.num_lstm_layers = num_lstm_layers
self.use_index_select = True
self.obj_dim = 5
self.col_dim = 3
self.con_dim = 2
self.goal_dim = flags.goal_dim
self.agent_loc_dim = 10
self.num_channels = (self.obj_dim + self.col_dim + self.con_dim + 1) * num_input_frames
if flags.disable_use_embedding:
print("not_using_embedding")
self.num_channels = (3+1+1+1+1)*num_input_frames
self.embed_object = nn.Embedding(11, self.obj_dim)
self.embed_color = nn.Embedding(6, self.col_dim)
self.embed_contains = nn.Embedding(4, self.con_dim)
self.embed_goal = nn.Embedding(self.observation_shape[0]*self.observation_shape[1] + 1, self.goal_dim)
self.embed_agent_loc = nn.Embedding(self.observation_shape[0]*self.observation_shape[1] + 1, self.agent_loc_dim)
init_ = lambda m: init(m, nn.init.orthogonal_, lambda x: nn.init.
constant_(x, 0), nn.init.calculate_gain('relu'))
self.feat_extract = nn.Sequential(
init_(nn.Conv2d(in_channels=self.num_channels, out_channels=32, kernel_size=(3, 3), stride=2, padding=1)),
nn.ELU(),
init_(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(3, 3), stride=2, padding=1)),
nn.ELU(),
init_(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(3, 3), stride=2, padding=1)),
nn.ELU(),
init_(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(3, 3), stride=2, padding=1)),
nn.ELU(),
init_(nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(3, 3), stride=2, padding=1)),
nn.ELU(),
)
self.fc = nn.Sequential(
init_(nn.Linear(32 + self.obj_dim + self.col_dim, self.state_embedding_dim)),
nn.ReLU(),
init_(nn.Linear(self.state_embedding_dim, self.state_embedding_dim)),
nn.ReLU(),
)
if use_lstm:
self.core = nn.LSTM(self.state_embedding_dim, self.state_embedding_dim, self.num_lstm_layers)
init_ = lambda m: init(m, nn.init.orthogonal_, lambda x: nn.init.
constant_(x, 0))
self.policy = init_(nn.Linear(self.state_embedding_dim, self.num_actions))
self.baseline = init_(nn.Linear(self.state_embedding_dim, 1))