| """GridPolicyMLP: a 500-32-32-5 tanh MLP policy network trained with PPO."""
|
| import torch
|
| import torch.nn as nn
|
|
|
|
|
| class GridPolicyMLP(nn.Module):
|
| """h1 = tanh(x @ W1 + b1); h2 = tanh(h1 @ W2 + b2); logits = h2 @ W3 + b3"""
|
|
|
| OBS_CHANNELS = ["agent", "moving objects", "reward targets", "obstacles", "previous frame"]
|
| ACTIONS = ["no-op", "up", "down", "left", "right"]
|
|
|
| def __init__(self):
|
| super().__init__()
|
| self.fc1 = nn.Linear(500, 32)
|
| self.fc2 = nn.Linear(32, 32)
|
| self.fc3 = nn.Linear(32, 5)
|
|
|
| def forward(self, x):
|
| h1 = torch.tanh(self.fc1(x))
|
| h2 = torch.tanh(self.fc2(h1))
|
| return self.fc3(h2)
|
|
|
| @torch.no_grad()
|
| def act(self, obs, deterministic=False):
|
| """obs: float array of 500 (or batch x 500) binary values. Returns action index/indices."""
|
| x = torch.as_tensor(obs, dtype=torch.float32)
|
| single = x.dim() == 1
|
| if single:
|
| x = x.unsqueeze(0)
|
| logits = self.forward(x)
|
| if deterministic:
|
| a = logits.argmax(dim=-1)
|
| else:
|
| a = torch.distributions.Categorical(logits=logits).sample()
|
| return int(a[0]) if single else a.tolist()
|
|
|
