| 246 | |
| 247 | |
| 248 | class Agent(nn.Module): |
| 249 | def __init__( |
| 250 | self, |
| 251 | num_actions, |
| 252 | hidden_size: int = 256, |
| 253 | c_v: int = 120, |
| 254 | c_k: int = 8, |
| 255 | c_s: int = 64, |
| 256 | num_queries: int = 4, |
| 257 | ): |
| 258 | """Agent implementing the attention agent. |
| 259 | """ |
| 260 | super(Agent, self).__init__() |
| 261 | self.hidden_size = hidden_size |
| 262 | self.c_v, self.c_k, self.c_s, self.num_queries = c_v, c_k, c_s, num_queries |
| 263 | |
| 264 | self.vision = VisionNetwork() |
| 265 | self.query = QueryNetwork() |
| 266 | # TODO: Implement SpatialBasis. |
| 267 | self.spatial = SpatialBasis() |
| 268 | |
| 269 | self.answer_processor = nn.Sequential( |
| 270 | # 1026 x 512 |
| 271 | nn.Linear( |
| 272 | (c_v + c_s) * num_queries + (c_k + c_s) * num_queries + 1 + 1, 512 |
| 273 | ), |
| 274 | nn.ReLU(), |
| 275 | nn.Linear(512, hidden_size), |
| 276 | ) |
| 277 | |
| 278 | self.policy_core = nn.LSTMCell(hidden_size, hidden_size) |
| 279 | |
| 280 | |
| 281 | self.prev_output = None |
| 282 | self.prev_hidden = None |
| 283 | |
| 284 | self.policy_head = nn.Sequential(nn.Linear(hidden_size, num_actions)) |
| 285 | self.values_head = nn.Sequential(nn.Linear(hidden_size, num_actions)) |
| 286 | |
| 287 | def reset(self): |
| 288 | self.vision.reset() |
| 289 | self.prev_output = None |
| 290 | self.prev_hidden = None |
| 291 | |
| 292 | def forward(self, X, prev_reward=None, prev_action=None): |
| 293 | |
| 294 | # 0. Setup. |
| 295 | # --------- |
| 296 | batch_size = X.size()[0] |
| 297 | if prev_reward is None: |
| 298 | # (n, 1, 1) |
| 299 | prev_reward = torch.stack([torch.zeros(1, 1)] * batch_size).to(X.device) |
| 300 | else: |
| 301 | prev_reward = prev_reward.unsqueeze(1).unsqueeze(1) |
| 302 | if prev_action is None: |
| 303 | # (n, 1, 1) |
| 304 | prev_action = torch.stack([torch.zeros(1, 1)] * batch_size).to(X.device) |
| 305 | else: |
nothing calls this directly
no outgoing calls
no test coverage detected